code: def setUp(self):
      super(DTensorMNISTTest, self).setUp()
      if config.list_physical_devices("GPU"):
          device_type = "GPU"
      elif test_util.is_tpu_present():
          device_type = "TPU"
      else:
          device_type = "CPU"
  
      local_devices = d_config.local_devices(device_type)
      num_devices = len(local_devices)
  
      global_device_ids = test_util.create_device_ids_array(
          (
              d_config.num_clients()
              * num_devices
              // _MODEL_DIM_SIZE.value,
              _MODEL_DIM_SIZE.value,
          )
      )
      device_ids_list = np.ravel(global_device_ids).tolist()
  
      index = d_config.client_id() * num_devices
      local_device_ids = device_ids_list[index : (index + num_devices)]
  
      self.mesh = d_layout.Mesh(
          [_BATCH_DIM, _MODEL_DIM],
          global_device_ids=global_device_ids,
          local_device_ids=local_device_ids,
          local_devices=local_devices,
      )
  
imports: import os
  from absl import flags
  import numpy
  from tensorflow.dtensor.python import accelerator_util
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import mesh_util
  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.eager import backprop
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import config
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  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 1 for calling the DTensorMNISTTest obj's underlying member method code with example usage: obj.setUp() without return types