code: def test_distribute_dataset_in_tf_function(self):
      strategy = mwms.MultiWorkerMirroredStrategy()
      local_batch_size = 4
      global_batch_size = (
          local_batch_size * strategy.num_replicas_in_sync
      )
      dataset = self.dataset.batch(global_batch_size).prefetch(2)
  
      distributed_dataset = strategy.experimental_distribute_dataset(
          dataset
      )
  
      @def_function.function
      def step_fn(iterator):
          images, labels = next(iterator)
          del labels
          return images
  
      result = strategy.run(step_fn, args=(iter(distributed_dataset),))
      self.assertIsInstance(
          result, dtensor_util.DTensorDistributedValue
      )
      self.assertLen(result.values, self.num_local_devices)
      for i in range(self.num_local_devices):
          self.assertEqual(
              result.values[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 7 for calling the StrategyDatasetTest obj's underlying member method code with example usage: obj.test_distribute_dataset_in_tf_function() and returns: images