class_docstr: Distribution for data parallelism.
  
  You can choose to create this instance by either specifing
  the `device_mesh` or `devices` arguments (but not both).
  
  The `device_mesh` argument is expected to be a `DeviceMesh` instance,
  and is expected to be 1D only. In case that the mesh has multiple axes,
  then the first axis will be treated as the data parallel dimension
  (and a warning will be raised).
  
  When a list of `devices` are provided, they will be used to construct a
  1D mesh.
  
  When both `mesh` and `devices` are absent, then `list_devices()`
  will be used to detect any available devices and create a 1D mesh from
  them.
  
  Args:
      device_mesh: Optional `DeviceMesh` instance.
      devices: Optional list of devices.
class_docstr_tok: ["Distribution","for","data","parallelism",".","You","can","choose","to","create","this","instance","by","either","specifing","the","`","device_mesh","`","or","`","devices","`","arguments","(","but","not","both",")",".","The","`","device_mesh","`","argument","is","expected","to","be","a","`","DeviceMesh","`","instance",",","and","is","expected","to","be","1D","only",".","In","case","that","the","mesh","has","multiple","axes",",","then","the","first","axis","will","be","treated","as","the","data","parallel","dimension","(","and","a","warning","will","be","raised",")",".","When","a","list","of","`","devices","`","are","provided",",","they","will","be","used","to","construct","a","1D","mesh",".","When","both","`","mesh","`","and","`","devices","`","are","absent",",","then","`","list_devices","(",")","`","will","be","used","to","detect","any","available","devices","and","create","a","1D","mesh","from","them",".","Args",":","device_mesh",":","Optional","`","DeviceMesh","`","instance",".","devices",":","Optional","list","of","devices","."]
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
  
      batch_size = tf_data_distribute.compute_batch_size(dataset)
      if 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)`"
          )
      per_worker_batch_size = tf_data_distribute.batch_sizes_for_worker(
          global_batch_size=batch_size,
          num_workers=self._num_process,
          num_replicas_per_worker=1,  # We hard code this for now.
          worker_index=self._process_id,
      )
      distributed_dataset = dataset.rebatch(per_worker_batch_size)
      distributed_dataset = tf_data_distribute._AutoShardDataset(
          distributed_dataset,
          num_workers=self._num_process,
          index=self._process_id,
          num_replicas=self._num_process,
      )
      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 8 for calling the DataParallel obj's underlying member method code with example usage: obj.distribute_dataset(dataset) and returns: distributed_dataset, dataset