class_docstr: Strategy extension that support both single and multi worker strategy.
code: def _distribute_datasets_from_function(self, dataset_fn, options):
      # TODO(scottzhu): Implement the logic for options in future
      del options
      input_context = distribute_lib.InputContext(
          num_input_pipelines=self._num_clients,
          input_pipeline_id=self._client_id,
          num_replicas_in_sync=self._num_replicas_in_sync,
      )
      dataset = dataset_fn(input_context)
  
      # Note that the dataset should already batched to local per-relica batch
      def _create_batch_layout(tensor_spec):
          rank = len(tensor_spec.shape)
          return layout.Layout.batch_sharded(
              self._mesh,
              batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
              rank=rank,
          )
  
      layouts = nest.map_structure(
          _create_batch_layout, dataset.element_spec
      )
  
      batch_size = distribute.compute_batch_size(dataset)
      # There are multiple case that the batch is not static, eg partial batch,
      # or uneven batch, in all those case, it will return -1.
      if batch_size.numpy() < 0:
          # When we don't have a static batch size.
          raise ValueError(
              "DTensor strategy requires a static batch size for now."
              "The dynamic batch size will be supported in future"
          )
      global_batch_size = (
          batch_size.numpy() * self._num_replicas_in_sync
      )
  
      return input_util.DTensorDataset(
          dataset=dataset,
          mesh=self._mesh,
          layouts=layouts,
          global_batch_size=global_batch_size,
          dataset_already_batched=True,
          batch_dim=dtensor_util.DEFAULT_BATCH_MESH_DIM_NAME,
          # TODO(scottzhu): Add prefetch support by inspecting the input dataset.
          prefetch=None,
          tf_data_service_config=None,
      )
  
imports: import functools
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.python.data.experimental.ops import distribute
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import distribute_utils
  from tensorflow.python.distribute.experimental import dtensor_util
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.util import nest
label_desc: Use image node_id 13 for calling the DTensorStrategyExtended obj's underlying member method code with example usage: obj._distribute_datasets_from_function(dataset_fn, options) and returns: input_util, layout