code: def testIdentityOpWithSparseTensorInputSimple(self, dtype):
      inputs = array_ops.ones([6, 4], dtype=dtype)
      layout = Layout.batch_sharded(self.mesh, _BATCH_DIM, rank=2)
  
      @polymorphic_function.function
      def f(x):
          return array_ops.identity(x)
  
      self.assertDTensorEqual(
          inputs,
          layout,
          f(numpy_util.pack_numpy(inputs, layout, make_sparse=True)),
      )
  
file_path: tensorflow-upstream/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 2 for calling the DTensorSPMDTest obj's underlying member method code with example usage: obj.testIdentityOpWithSparseTensorInputSimple(dtype) and returns: array_ops