code: def testIdentityOpWithSparseTensorInputComplex(
      self, dtype, is_sparse_a, is_sparse_b
  ):
      inputs_a = array_ops.ones([2, 1], dtype=dtype)
      inputs_b = array_ops.ones([32, 16], dtype=dtype)
  
      layout_a = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
      layout_b = Layout.replicated(self.mesh, rank=2)
  
      @polymorphic_function.function
      def f(x, y):
          return array_ops.identity(x), array_ops.identity(y)
  
      got_a, got_b = f(
          numpy_util.pack_numpy(
              inputs_a, layout_a, make_sparse=is_sparse_a
          ),
          numpy_util.pack_numpy(
              inputs_b, layout_b, make_sparse=is_sparse_b
          ),
      )
  
      self.assertDTensorEqual(inputs_a, layout_a, got_a)
      self.assertDTensorEqual(inputs_b, layout_b, got_b)
  
file_path: tensorflow/tensorflow/dtensor/python/tests/sparse_test.py
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import layout
  from tensorflow.dtensor.python import numpy_util
  from tensorflow.dtensor.python.tests import test_util
  from tensorflow.python.eager.polymorphic_function import polymorphic_function
  from tensorflow.python.framework import dtypes
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 3 for calling the DTensorSPMDTest obj's underlying member method code with example usage: obj.testIdentityOpWithSparseTensorInputComplex(dtype, is_sparse_a, is_sparse_b) and returns: array_ops,  array_ops