code: def testRangeIteration(self, is_graph):
      batch_size = 8
      num_batches = 4
      dataset = dataset_ops.DatasetV2.from_tensor_slices(
          self._images([batch_size * num_batches, 8, 8, 3])
      )
      images_layout = Layout.batch_sharded(
          self.mesh, batch_dim=MESH_DIM_BATCH, rank=4
      )
  
      d_dataset = input_util.DTensorDataset(
          dataset=dataset,
          global_batch_size=batch_size,
          mesh=self.mesh,
          layouts=images_layout,
          batch_dim=MESH_DIM_BATCH,
      )
  
      def train(iterator, steps):
          iters = 1
          output = next(iterator)
          for _ in math_ops.range(steps - 1):
              output += next(iterator)
              iters += 1
              if not is_graph:
                  mesh_util.barrier(self.mesh)
          return output, iters
  
      train_fn = (
          polymorphic_function.function(train) if is_graph else train
      )
      exception = (
          errors_impl.OutOfRangeError if is_graph else StopIteration
      )
  
      iterator = iter(dataset.batch(batch_size, drop_remainder=True))
      output, iters = train_fn(iterator, num_batches)
  
      d_iterator = iter(d_dataset)
      d_output, d_iters = train_fn(d_iterator, num_batches)
  
      mesh_util.barrier(self.mesh)
      # Try one more iteration which will raise an exception since the iterator is
      # exhausted.
      with self.assertRaises(exception):
          if is_graph:
              # FIXME(b/285884302): This flakily raises error
              # "Cannot add 'while_cond' function, because a different function"
              # Since num_batches is changed to 1, it retriggers SPMD expansion.
              # Recreating polymorphic function to avoid running into the error.
              train_fn = polymorphic_function.function(train)
          train_fn(d_iterator, 1)
          # In the graph case, we need to wait for the executor to finish all async
          # calls after invoking the tf.function to ensure any pending error is
          # raised.
          mesh_util.barrier(self.mesh)
  
      self.assertEqual(iters, d_iters)
      self.assertDTensorEqual(output, images_layout, d_output)
  
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 6 for calling the DTensorDatasetTest obj's underlying member method code with example usage: obj.testRangeIteration(is_graph) and returns: output,  iters