code: def testMnist(self, on_function):
      mnist_dataset = mnist_fake_dataset()
  
      (n_w, n_b, n_k), (w, b, k) = self.init_var(self.mesh)
  
      n_dataset = mnist_dataset.batch(_BATCH_SIZE, drop_remainder=True)
      n_iter = iter(n_dataset)
  
      input_layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=4)
      label_layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
      dtensor_dataset = input_util.DTensorDataset(
          dataset=mnist_dataset,
          global_batch_size=_BATCH_SIZE,
          mesh=self.mesh,
          layouts=(input_layout, label_layout),
          batch_dim=_BATCH_DIM,
      )
      dtensor_iter = iter(dtensor_dataset)
  
      step_fn = (
          polymorphic_function.function(_run_step)
          if on_function
          else _run_step
      )
  
      # Training loop.
      for _ in range(_STEPS):
          # Normal run without DTensor.
          n_input, _ = next(n_iter)
          g_nw, g_nb, n_loss = step_fn(n_input, n_w, n_b, n_k)
  
          # DTensor Run
          dtensor_input, _ = next(dtensor_iter)
          with ops.device_v2(api.device_name()):
              gw, gb, loss = step_fn(dtensor_input, w, b, k)
  
          loss_unpack = api.unpack(loss)
          self.assertAllEqual(loss_unpack[0], loss_unpack[1])
  
          self.assertAllClose(n_loss, loss, atol=_ATOL, rtol=_RTOL)
          self.assertAllClose(g_nw, gw, atol=_ATOL, rtol=_RTOL)
          self.assertAllClose(g_nb, gb, atol=_ATOL, rtol=_RTOL)
          self.assertAllClose(n_w, w, atol=_ATOL, rtol=_RTOL)
          self.assertAllClose(n_b, b, atol=_ATOL, rtol=_RTOL)
  
imports: 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 input_util
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.data.ops import dataset_ops
  from tensorflow.python.eager import backprop
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import array_ops_stack
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_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 3 for calling the DTensorMNISTTest obj's underlying member method code with example usage: obj.testMnist(on_function) without return types