keras/src/distribution/distribution_lib.py/ModelParallel/distribute_dataset
class ModelParallel:    def distribute_dataset(self, dataset):
        from tensorflow.python.data.experimental.ops import (
            distribute as tf_data_distribute,
        )

        from keras.src.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 across 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 dimension.
            return dataset.prefetch(tf.data.AUTOTUNE)
        else:
            if mesh_batch_dim_size % local_device_count != 0:
                raise ValueError(
                    "The Batch dimension 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)