class_docstr: Distribution that shards model variables.
  
  Compare to `DataParallel` which replicates the variables across all devices,
  `ModelParallel` allows you to shard variables in addition to the input data.
  
  To construct a `ModelParallel` distribution, you need to provide a
  `DeviceMesh` and a `LayoutMap`.
  
  1. `DeviceMesh` contains physcial device information. The axis names in
      the mesh will be used to map the variable and data layout.
  2. `LayoutMap` contains the mapping between variable paths to their
      corresponding `TensorLayout`.
  
  Example:
  
  ```python
  devices = list_devices()    # Assume there are 8 devices.
  
  # Create a mesh with 2 devices for data parallelism and 4 devices for
  # model parallelism.
  device_mesh = DeviceMesh(shape=(2, 4), axis_names=('batch', 'model'),
                           devices=devices)
  # Create a layout map that shard the `Dense` layer and `Conv2D`
  # layer variables on the last dimension.
  # Based on the `device_mesh`, this means the variables
  # will be split across 4 devices. Any other variable that doesn't
  # match any key in the layout map will be fully replicated.
  layout_map = LayoutMap(device_mesh)
  layout_map['dense.*kernel'] = (None, 'model')
  layout_map['dense.*bias'] = ('model',)
  layout_map['conv2d.*kernel'] = (None, None, None, 'model')
  layout_map['conv2d.*bias'] = ('model',)
  
  distribution = ModelParallel(device_mesh=device_mesh,
                               layout_map=layout_map,
                               batch_dim_name='batch')
  # Set the global distribution, or via `with distribution.scope():`
  set_distribution(distribution)
  
  model = model_creation()
  model.compile()
  model.fit(data)
  ```
  
  You can quickly update the device mesh shape to change the sharding factor
  of the variables. E.g.
  ```
  # With only the shape change for the device mesh, the variables will be
  # sharded across 8 devices instead of 4, which further reduces the memory
  # footprint of variables on each of the device.
  device_mesh = DeviceMesh(shape=(1, 8), axis_names=('batch', 'model'),
                           devices=devices)
  ```
  
  To figure out a proper layout mapping rule for all the model variables, you
  can first list out all the model variable paths, which will be used as the
  key to map the variables to `TensorLayout`.
  
  e.g.
  ```
  model = create_model()
  for v in model.variables:
      print(v.path)
  ```
  
  Args:
      device_mesh: `DeviceMesh` instance for physical device and its
          logical mapping.
      layout_map: `LayoutMap` instance which map the variable path to the
          corresponding `TensorLayout`. The axis names of the
          `TensorLayout`s should match to the axis names in the
          device_mesh, or exception will be raised.
      batch_dim_name: optional string, the axis name in the `device_mesh`
          that will be used to distribute data. If unspecified, the
          first axis from the `device_mesh` will be used.
class_docstr_tok: ["Distribution","that","shards","model","variables",".","Compare","to","`","DataParallel","`","which","replicates","the","variables","across","all","devices",",","`","ModelParallel","`","allows","you","to","shard","variables","in","addition","to","the","input","data",".","To","construct","a","`","ModelParallel","`","distribution",",","you","need","to","provide","a","`","DeviceMesh","`","and","a","`","LayoutMap","`",".","1",".","`","DeviceMesh","`","contains","physcial","device","information",".","The","axis","names","in","the","mesh","will","be","used","to","map","the","variable","and","data","layout",".","2",".","`","LayoutMap","`","contains","the","mapping","between","variable","paths","to","their","corresponding","`","TensorLayout","`",".","Example",":","``","`","python","devices","=","list_devices","(",")","#","Assume","there","are","8","devices",".","#","Create","a","mesh","with","2","devices","for","data","parallelism","and","4","devices","for","#","model","parallelism",".","device_mesh","=","DeviceMesh","(","shape=","(","2",",","4",")",",","axis_names=","(","'batch","'",",","'model","'",")",",","devices=devices",")","#","Create","a","layout","map","that","shard","the","`","Dense","`","layer","and","`","Conv2D","`","#","layer","variables","on","the","last","dimension",".","#","Based","on","the","`","device_mesh","`",",","this","means","the","variables","#","will","be","split","across","4","devices",".","Any","other","variable","that","doesn't","#","match","any","key","in","the","layout","map","will","be","fully","replicated",".","layout_map","=","LayoutMap","(","device_mesh",")","layout_map","[","'dense",".","*","kernel","'","]","=","(","None",",","'model","'",")","layout_map","[","'dense",".","*","bias","'","]","=","(","'model","'",",",")","layout_map","[","'conv2d",".","*","kernel","'","]","=","(","None",",","None",",","None",",","'model","'",")","layout_map","[","'conv2d",".","*","bias","'","]","=","(","'model","'",",",")","distribution","=","ModelParallel","(","device_mesh=device_mesh",",","layout_map=layout_map",",","batch_dim_name='batch","'",")","#","Set","the","global","distribution",",","or","via","`","with","distribution.scope","(",")",":","`","set_distribution","(","distribution",")","model","=","model_creation","(",")","model.compile","(",")","model.fit","(","data",")","``","`","You","can","quickly","update","the","device","mesh","shape","to","change","the","sharding","factor","of","the","variables",".","E.g",".","``","`","#","With","only","the","shape","change","for","the","device","mesh",",","the","variables","will","be","#","sharded","across","8","devices","instead","of","4",",","which","further","reduces","the","memory","#","footprint","of","variables","on","each","of","the","device",".","device_mesh","=","DeviceMesh","(","shape=","(","1",",","8",")",",","axis_names=","(","'batch","'",",","'model","'",")",",","devices=devices",")","``","`","To","figure","out","a","proper","layout","mapping","rule","for","all","the","model","variables",",","you","can","first","list","out","all","the","model","variable","paths",",","which","will","be","used","as","the","key","to","map","the","variables","to","`","TensorLayout","`",".","e.g",".","``","`","model","=","create_model","(",")","for","v","in","model.variables",":","print","(","v.path",")","``","`","Args",":","device_mesh",":","`","DeviceMesh","`","instance","for","physical","device","and","its","logical","mapping",".","layout_map",":","`","LayoutMap","`","instance","which","map","the","variable","path","to","the","corresponding","`","TensorLayout","`",".","The","axis","names","of","the","`","TensorLayout","`","s","should","match","to","the","axis","names","in","the","device_mesh",",","or","exception","will","be","raised",".","batch_dim_name",":","optional","string",",","the","axis","name","in","the","`","device_mesh","`","that","will","be","used","to","distribute","data",".","If","unspecified",",","the","first","axis","from","the","`","device_mesh","`","will","be","used","."]
code: def distribute_dataset(self, dataset):
      from tensorflow.python.data.experimental.ops import (
          distribute as tf_data_distribute,
      )
  
      from keras.utils.module_utils import tensorflow as tf
  
      if not isinstance(dataset, tf.data.Dataset):
          raise ValueError(
              "Only `tf.data.Dataset` is supported for "
              f"sharding, got {type(dataset)}"
          )
      if not self._is_multi_process:
          return dataset
  
      global_batch_size = tf_data_distribute.compute_batch_size(dataset)
      if global_batch_size.numpy() < 0:
          raise ValueError(
              "The batch size of the input dataset is "
              "unknown. Please config the batch size for "
              "the input dataset, e.g via `dataset.batch(batch_size)`"
          )
  
      # We need to compute the the per-worker batch size.
      # Imagine we have 4 workers, and each worker has 2 devices. The device
      # ID will be w{worker_id}_{device_id} from global perspective, eg
      # 'w0_0', 'w0_1', 'w1_0', etc.
      # For a mesh with global shape {'batch': 4, 'model': 2}, we would like
      # the mesh to be like below:
      #         batch ->
      #         --------------------------
      # model   | w0_0, w0_1, w1_0, w1_1 |
      #   v     | w0_2, w0_3, w1_2, w1_3 |
      #         --------------------------
      # In this case, the batch dim will be split acorss 2 workers.
      #
      # In the case that global batch dim is smaller than the worker size, eg
      # {'batch': 1, 'model': 8}, then the mesh will be like below:
      #         batch ->
      #         --------
      # model   | w0_0 |
      #         | w0_1 |
      #         | w1_0 |
      #         | w1_1 |
      #         | w2_0 |
      #         | w2_1 |
      #         | w3_0 |
      #         | w3_1 |
      #         --------
      # And the data batch will be fully replicated for each of the worker.
      mesh_batch_dim_index = self.device_mesh.axis_names.index(
          self._batch_dim_name
      )
      mesh_batch_dim_size = self.device_mesh.shape[mesh_batch_dim_index]
      # TODO(scottzhu): local device count can be a field of the mesh.
      local_device_count = (
          np.prod(self.device_mesh.shape) // self._num_process
      )
      if mesh_batch_dim_size < local_device_count:
          # No sharding is needed in this case. The worker will have the
          # global batch size, and data from the iterator will need to be
          # replicated for model dimention.
          return dataset.prefetch(tf.data.AUTOTUNE)
      else:
          if mesh_batch_dim_size % local_device_count != 0:
              raise ValueError(
                  "The Batch dimention of the mesh is not compatible "
                  "with the local worker device count. Mesh batch "
                  f"dim = {mesh_batch_dim_size} and local device "
                  f"count = {local_device_count}"
              )
          num_shards = mesh_batch_dim_size // local_device_count
          per_worker_batch_size = global_batch_size // num_shards
  
          distributed_dataset = dataset.rebatch(per_worker_batch_size)
          distributed_dataset = tf_data_distribute._AutoShardDataset(
              distributed_dataset,
              num_workers=num_shards,
              index=self._process_id % num_shards,
              num_replicas=num_shards,
          )
          return distributed_dataset.prefetch(tf.data.AUTOTUNE)
  
imports: import collections
  import contextlib
  import os
  import re
  import warnings
  import numpy
  from keras.api_export import keras_export
  from keras.backend import KerasTensor
  from keras.backend import distribution_lib
  from keras.backend.common import global_state
label_desc: Use image node_id 5 for calling the ModelParallel obj's underlying member method code with example usage: obj.distribute_dataset(dataset) and returns: dataset, dataset, distributed_dataset