code: def testIterOnBatchedDataset(self, is_graph):
      dataset = dataset_ops.DatasetV2.from_tensors(
          {
              "images": self.images,
              "labels": self.labels,
          }
      ).repeat()
  
      images_layout = Layout.batch_sharded(
          self.mesh, batch_dim=MESH_DIM_BATCH, rank=4
      )
      labels_layout = Layout.batch_sharded(
          self.mesh, batch_dim=MESH_DIM_BATCH, rank=2
      )
      layouts = {"images": images_layout, "labels": labels_layout}
  
      global_batch_size = 32
      per_replica_batch_size = global_batch_size // MESH_SIZE_BATCH
      batched_dataset = dataset.batch(
          per_replica_batch_size, drop_remainder=True
      )
  
      d_dataset = input_util.DTensorDataset(
          dataset=batched_dataset,
          global_batch_size=global_batch_size,
          dataset_already_batched=True,
          mesh=self.mesh,
          layouts=layouts,
          batch_dim=MESH_DIM_BATCH,
      )
  
      def train(iterator):
          return next(iterator)
  
      train_fn = (
          polymorphic_function.function(train) if is_graph else train
      )
  
      d_iterator = iter(d_dataset)
      d_element = train_fn(d_iterator)
  
      expected = next(
          iter(dataset.batch(global_batch_size, drop_remainder=True))
      )
      self.assertDTensorEqual(
          expected["images"], images_layout, d_element["images"]
      )
      self.assertDTensorEqual(
          expected["labels"], labels_layout, d_element["labels"]
      )
  
imports: import contextlib
  import threading
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import tensor_shape
  from tensorflow.python.framework import tensor_spec
  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 11 for calling the DTensorDatasetTest obj's underlying member method code with example usage: obj.testIterOnBatchedDataset(is_graph) and returns: next