code: def testSparseTensorDenseMatMul(
      self, is_sparse_a, is_sparse_b, shard_type
  ):
      inputs_a = array_ops.ones([16, 16])
      inputs_b = array_ops.ones([16, 16])
  
      if shard_type == "Replicated":
          layout_a = Layout.replicated(self.mesh, rank=2)
          layout_b = Layout.replicated(self.mesh, rank=2)
      else:
          layout_a = Layout([_MESH_DIM_X, _BATCH_DIM], self.mesh)
          layout_b = Layout(["unsharded", _MESH_DIM_X], self.mesh)
  
      expected = math_ops.matmul(inputs_a, inputs_b)
  
      @polymorphic_function.function
      def f(x, y):
          return math_ops.matmul(x, y)
  
      got = 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(
          expected, Layout.replicated(self.mesh, rank=2), got
      )
  
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 5 for calling the DTensorSPMDTest obj's underlying member method code with example usage: obj.testSparseTensorDenseMatMul(is_sparse_a, is_sparse_b, shard_type) and returns: math_ops