code: def testSparseSoftmaxCrossentropyWithLogits(
      self, labels_layout, logits_layout
  ):
      expected_layout = Layout.replicated(self.mesh, rank=1)
      if labels_layout == 1 or logits_layout == 1 or logits_layout == 3:
          expected_layout = Layout.inner_sharded(
              self.mesh, _MESH_DIM_X, rank=1
          )
      elif labels_layout == 2:
          expected_layout = Layout.inner_sharded(
              self.mesh, _MESH_DIM_Y, rank=1
          )
  
      labels_layout = self.layouts_1d[labels_layout]
      logits_layout = self.layouts_2d[logits_layout]
      labels_numpy = np.random.randint(size=[6], low=0, high=4)
      logits_numpy = np.random.uniform(size=[6, 4])
      labels = constant_op.constant(labels_numpy, dtype=dtypes.int64)
      logits = constant_op.constant(logits_numpy, dtype=dtypes.float32)
  
      # Should we test against the built in version or the patched version?
      expected = nn_ops.sparse_softmax_cross_entropy_with_logits_v2(
          labels=labels, logits=logits
      )
  
      labels = numpy_util.pack_numpy(labels, labels_layout)
      logits = numpy_util.pack_numpy(logits, logits_layout)
      dtensor_result = (
          nn_ops.sparse_softmax_cross_entropy_with_logits_v2(
              labels=labels, logits=logits
          )
      )
      self.assertDTensorEqual(expected, expected_layout, dtensor_result)
  
imports: from absl.testing import parameterized
  import numpy
  from tensorflow.dtensor.python import api
  from tensorflow.dtensor.python import d_random
  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 constant_op
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import errors_impl
  from tensorflow.python.framework import ops
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import stateless_random_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 6 for calling the APITest obj's underlying member method code with example usage: obj.testSparseSoftmaxCrossentropyWithLogits(labels_layout, logits_layout) without return types