code: def testStatelessRNGOpsWithSingleDimensionSharded(
      self, op, dtype, op_version, dim, shard_type
  ):
      shape = [128, 128]
      seed = [123, 321]
      sharding = self.shardings[shard_type]
      layout = sharding(self.mesh, dim, rank=2)
  
      # Raw rng Ops do not have inputs, so we need to place the Op DTensor device
      # explicitly.
      with ops.device_v2(api.device_name()):
          with api._dtensor_device()._default_layout(layout):
              b = _call_dtensor_op(
                  op=op,
                  seed=seed,
                  shape=shape,
                  dtype=dtype,
                  key=self.key,
                  counter=self.counter,
                  alg=self.alg,
                  minval=self.minval,
                  maxval=self.maxval,
                  op_version=op_version,
                  mesh=self.mesh,
              )
  
      api.check_layout(b, layout)
      b = [tensor.numpy() for tensor in api.unpack(b)]
  
      if dim == _MESH_DIM_X:
          if shard_type == "batch":
              self.assertAllEqual(b[0].shape, [64, 128])
          else:
              assert shard_type == "inner"
              self.assertAllEqual(b[0].shape, [128, 64])
  
          # first check that each component is same as the row header.
          for i in range(self.mesh.num_local_devices()):
              self.assertAllEqual(b[i], b[_ROW_HEAD[i]])
          # then check the row header are NOT identital.
          self.assertNotAllEqual(b[_ROW_0_HEAD], b[_ROW_1_HEAD])
  
      elif dim == _MESH_DIM_Y:
          if shard_type == "batch":
              self.assertAllEqual(b[0].shape, [32, 128])
          else:
              assert shard_type == "inner"
              self.assertAllEqual(b[0].shape, [128, 32])
  
          # first check elements in same columns are identical
          for i in range(self.mesh.num_local_devices()):
              self.assertAllEqual(b[i], b[_COL_HEAD[i]])
  
          col_heads = [
              _COL_0_HEAD,
              _COL_1_HEAD,
              _COL_2_HEAD,
              _COL_3_HEAD,
          ]
          # then check the column header are not identital (mutually)
          for i in range(self.mesh.num_local_devices() - 1):
              for j in range(self.mesh.num_local_devices()):
                  if i == j:
                      continue
                  if i in col_heads and j in col_heads:
                      self.assertNotAllEqual(b[i], b[j])
  
      else:
          self.fail("should not reach here.")
  
imports: from absl.testing import parameterized
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.dtensor.python.tests import test_util_ops
  from tensorflow.python.distribute import tpu_strategy
  from tensorflow.python.distribute.cluster_resolver.tpu import tpu_cluster_resolver
  from tensorflow.python.eager import remote
  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 ops
  from tensorflow.python.ops import gen_bitwise_ops
  from tensorflow.python.ops import gen_stateful_random_ops
  from tensorflow.python.ops import gen_stateless_random_ops_v2
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import variables
  from tensorflow.python.platform import test
  from tensorflow.python.tpu import device_assignment
label_desc: Use image node_id 4 for calling the DTensorRNGTest obj's underlying member method code with example usage: obj.testStatelessRNGOpsWithSingleDimensionSharded(op, dtype, op_version, dim, shard_type) without return types