code: def testCollectiveOpsOnComplex64(self):
      # This functions tests for AllScatter, AllGather, and AllReduce.
      a = constant_op.constant(
          np.array([[1, 2 + 2j], [3 + 1j, 4 + 5j]]),
          dtype=dtypes.complex64,
      )
      # Tests AllScatter
      sharded_a = api.relayout(
          a, self.first_dimension_sharded_layout_2d
      )
      # Tests AllGather / AllReduce
      unsharded_a = api.relayout(
          sharded_a, self.fully_replicated_layout_2d
      )
  
      self.assertDTensorEqual(
          a, self.fully_replicated_layout_2d, unsharded_a
      )
  
imports: import os
  from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import config
  from tensorflow.dtensor.python import d_variable
  from tensorflow.dtensor.python import dtensor_device
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python.tests import test_backend_util
  from tensorflow.dtensor.python.tests import test_util
  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.ops import array_ops
  from tensorflow.python.ops import gen_array_ops
  from tensorflow.python.ops import gen_math_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 9 for calling the CollectiveTest obj's underlying member method code with example usage: obj.testCollectiveOpsOnComplex64() without return types