code: def testNthSmallestReduceSecondDim(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      a_dim = mtf.Dimension("a", 6)
      b_dim = mtf.Dimension("b", 2)
      inputs = tf.constant(
          [[1, 10], [2, 9], [3, 8], [4, 7], [5, 6], [6, 5]]
      )
      n = 0  # find smallest element (n is zero-indexed)
      reduced_dim = b_dim
      expected_outputs = tf.constant([1, 2, 3, 4, 5, 5])
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([a_dim, b_dim])
      )
      mtf_outputs = mtf.nth_smallest_element(
          mtf_inputs, n, reduced_dim, "test_nth_smallest"
      )
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          shape="all:2", layout="a:all", devices=["", ""]
      )
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
      self.assertAllEqual(
          self.evaluate(actual_outputs), self.evaluate(expected_outputs)
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import mesh_tensorflow
  from mesh_tensorflow import test_utils
  import numpy
  import tensorflow.compat.v1
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 2 for calling the NthSmallestTest obj's underlying member method code with example usage: obj.testNthSmallestReduceSecondDim() without return types