keras/src/distribution/distribution_lib.py/ModelParallel/get_data_layout
class ModelParallel:    def get_data_layout(self, data_shape):
        data_shard_spec = [None] * len(data_shape)
        data_shard_spec[0] = self._batch_dim_name  # Shard on the first dim
        return TensorLayout(data_shard_spec, self.device_mesh)