code: def test_create_batched_dataset(self):
      strategy = mwms.MultiWorkerMirroredStrategy()
      global_batch_size = self.num_client * self.num_local_devices * 2
      dataset = self.dataset.batch(global_batch_size).prefetch(2)
  
      distributed_dataset = strategy.experimental_distribute_dataset(
          dataset
      )
      element = next(iter(distributed_dataset))
      batched_image, batched_label = element
      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.
      self.assertLen(
          d_api.unpack(batched_image), self.num_local_devices
      )
      self.assertLen(
          d_api.unpack(batched_label), self.num_local_devices
      )
  
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 2 for calling the StrategyDatasetTest obj's underlying member method code with example usage: obj.test_create_batched_dataset() without return types