code: def testCholesky(self, num_batch_dim):
      # Input needs to be symmetric and positive definite.
      x = constant_op.constant(
          [[1, 1, 1, 1], [1, 5, 5, 5], [1, 5, 14, 14], [1, 5, 14, 17]],
          dtype=dtypes.float32,
      )
      for _ in range(num_batch_dim):
          x = array_ops.expand_dims_v2(x, 0)
          s = [4] + [1 for _ in range(array_ops.rank(x) - 1)]
          x = gen_array_ops.tile(x, s)
  
      expected_result = gen_linalg_ops.cholesky(x)
  
      if num_batch_dim == 0:
          layout_spec = []
      elif num_batch_dim == 1:
          layout_spec = [self._mesh_dim_b]
      elif num_batch_dim == 2:
          layout_spec = [self._mesh_dim_b, self._mesh_dim_x]
      layout = Layout(layout_spec + ["unsharded"] * 2, self.mesh)
  
      x = numpy_util.pack_numpy(x, layout)
      got = gen_linalg_ops.cholesky(input=x)
      self.assertDTensorEqual(expected_result, layout, got)
  
imports: import itertools
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  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.eager import backprop
  from tensorflow.python.eager import context
  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 array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_image_ops
  from tensorflow.python.ops import gen_linalg_ops
  from tensorflow.python.ops import nn_impl
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 2 for calling the DTensorBatchParallelSPMDTest obj's underlying member method code with example usage: obj.testCholesky(num_batch_dim) without return types