code: def testIterWithLayouts(
      self, images_sharding, labels_sharding, is_graph, through_dtensor
  ):
      if through_dtensor:
          scope = api.default_mesh(self.mesh)
      else:
          scope = contextlib.nullcontext()
  
      with scope:
          batch_size = 32
          dataset = dataset_ops.DatasetV2.from_tensors(
              (self.images, self.labels)
          ).repeat()
          batched_dataset = dataset.batch(
              batch_size, drop_remainder=True
          )
  
          images_layout = Layout(images_sharding, self.mesh)
          labels_layout = Layout(labels_sharding, self.mesh)
          layouts = (images_layout, labels_layout)
          batch_dim = None
          if (
              MESH_DIM_BATCH in images_sharding
              or MESH_DIM_BATCH in labels_sharding
          ):
              batch_dim = MESH_DIM_BATCH
  
          d_dataset = input_util.DTensorDataset(
              dataset=dataset,
              global_batch_size=batch_size,
              mesh=self.mesh,
              layouts=layouts,
              batch_dim=batch_dim,
          )
  
          def train(iterator):
              return next(iterator)
  
          train_fn = (
              polymorphic_function.function(train)
              if is_graph
              else train
          )
  
          d_iterator = iter(d_dataset)
          d_images, d_labels = train_fn(d_iterator)
  
      iterator = iter(batched_dataset)
      images, labels = train_fn(iterator)
  
      self.assertDTensorEqual(images, images_layout, d_images)
      self.assertDTensorEqual(labels, labels_layout, d_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 15 for calling the DTensorDatasetTest obj's underlying member method code with example usage: obj.testIterWithLayouts(images_sharding, labels_sharding, is_graph, through_dtensor) and returns: next