code: def testVariadicSortLexicographic(self, key_type_2):
      # Three inputs: the first two are used for lexicographic sort, and the
      # third is just swapped accordingly.
      # The first array will contain only 0 and 1, to test lexicographic order
      if np.__version__ < "1.15":
          raise unittest.SkipTest(
              "np.take_along_axis was added in 1.15"
          )
      shape = (20,)
      if key_type_2 not in self._supported_key_types():
          return
      for key_type_1 in [np.int16, np.uint16, np.int32, np.uint32]:
          for value_type in self._supported_key_types():
              inputs = [
                  # Ensure that some keys in the first input are equal
                  np.random.uniform(0, 2, shape).astype(key_type_1),
                  self._shuffled_arange(shape, key_type_2),
                  self._shuffled_arange(shape, value_type),
              ]
              # The first two arrays are sorted lexicographically, and the third
              # is shuffled the same way
              sorted_indices = np.argsort(100 * inputs[0] + inputs[1])
              expected = [
                  np.take_along_axis(inp, sorted_indices, axis=0)
                  for inp in inputs
              ]
  
              @function.Defun(
                  key_type_1,
                  key_type_1,
                  key_type_2,
                  key_type_2,
                  value_type,
                  value_type,
              )
              def compare_lexicographic(x1, x2, y1, y2, z1, z2):
                  del z1, z2
                  return math_ops.logical_or(
                      x1 < x2,
                      math_ops.logical_and(
                          math_ops.equal(x1, x2), y1 < y2
                      ),
                  )
  
              def wrap_sort(*args):
                  return xla.variadic_sort(
                      args,  # Pass the arguments as a tuple
                      comparator=compare_lexicographic,
                      dimension=0,
                      is_stable=False,
                  )
  
              self._assertOpOutputMatchesExpected(
                  wrap_sort, inputs, expected=expected
              )
  
imports: import unittest
  from absl.testing import parameterized
  import numpy
  from tensorflow.compiler.tests import xla_test
  from tensorflow.compiler.tf2xla.python import xla
  from tensorflow.python.framework import dtypes
  from tensorflow.python.framework import function
  from tensorflow.python.framework import tensor
  from tensorflow.python.framework import test_util
  from tensorflow.python.ops import array_ops
  from tensorflow.python.ops import math_ops
  from tensorflow.python.ops import nn_ops
  from tensorflow.python.ops import sort_ops
  from tensorflow.python.platform import test
label_desc: Use image node_id 9 for calling the XlaSortOpTest obj's underlying member method code with example usage: obj.testVariadicSortLexicographic(key_type_2) and returns: math_ops, xla