code: def test_distribute_dataset_from_fn(self):
      strategy = mwms.MultiWorkerMirroredStrategy()
      local_batch_size = 4
      global_batch_size = (
          local_batch_size * strategy.num_replicas_in_sync
      )
  
      def dataset_fn(option):
          del option
          return (
              dataset_ops.Dataset.from_tensors(
                  (self.images, self.labels)
              )
              .repeat()
              .batch(local_batch_size, drop_remainder=True)
              .prefetch(2)
          )
  
      distributed_dataset = strategy.distribute_datasets_from_function(
          dataset_fn, None
      )
      iterator = iter(distributed_dataset)
  
      self.assertEqual(
          distributed_dataset.element_spec,
          (
              tensor_spec.TensorSpec(
                  shape=(global_batch_size, 8, 8, 3),
                  dtype=dtypes.float32,
                  name=None,
              ),
              tensor_spec.TensorSpec(
                  shape=(global_batch_size, 1),
                  dtype=dtypes.float32,
                  name=None,
              ),
          ),
      )
      self.assertEqual(
          distributed_dataset.element_spec, iterator.element_spec
      )
  
      batched_image, batched_label = next(iterator)
      self.assertEqual(
          batched_image.shape, [global_batch_size, 8, 8, 3]
      )
      self.assertEqual(batched_label.shape, [global_batch_size, 1])
  
      # After unpack, it should only get the local shards.
      unpacked_images = d_api.unpack(batched_image)
      self.assertLen(unpacked_images, self.num_local_devices)
      for i in range(self.num_local_devices):
          self.assertEqual(
              unpacked_images[i].shape, [local_batch_size, 8, 8, 3]
          )
  
imports: import json
  import os
  from absl import flags
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import multi_client_test_util
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.distribute import distribute_lib
  from tensorflow.python.distribute import reduce_util
  from tensorflow.python.distribute.cluster_resolver import tfconfig_cluster_resolver
  from tensorflow.python.distribute.experimental import dtensor_util
  from tensorflow.python.distribute.experimental import multi_worker_mirrored_strategy
  from tensorflow.python.eager import def_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import tensor_spec
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
label_desc: Use image node_id 5 for calling the StrategyDatasetTest obj's underlying member method code with example usage: obj.test_distribute_dataset_from_fn() and returns: dataset_ops