code: def testTopK(self, dtype, rank):
      if dtype in self.numeric_types:
          # Use small input size for bfloat16. Otherwise, we'll get duplicate values
          # after conversion to bfloat16, so the possible resulting index array is
          # no longer unique.
          if dtype in (dtypes.bfloat16.as_numpy_dtype, np.float16):
              array_size = 20
              k_options = [0, 1, 2, 10, 20]
          elif dtype in (
              dtypes.uint8.as_numpy_dtype,
              dtypes.int8.as_numpy_dtype,
          ):
              array_size = 111
              k_options = [0, 1, 2, 10, 20]
          else:
              array_size = 200 * 1000
              k_options = [0, 1, 2, 10, 20, 100, 1000, 200 * 1000]
  
          # Tile array to tensor of specified rank, then shuffle along the last dim
          x = np.arange(array_size)
          x = np.tile(x, (2,) * (rank - 1) + (1,))
          np.apply_along_axis(np.random.shuffle, -1, x)
  
          sorted_indices = x.argsort(axis=-1)[..., ::-1]
          sorted_values = np.sort(x, axis=-1)[..., ::-1]
          for k in k_options:
              indices = sorted_indices[..., :k]
              expected = sorted_values[..., :k]
  
              def topk(v, k=k):
                  return nn_ops.top_k(v, k=k, sorted=True)
  
              self._assertOpOutputMatchesExpected(
                  topk,
                  [x.astype(dtype)],
                  expected=[expected.astype(dtype), indices],
              )
  
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 12 for calling the XlaSortOpTest obj's underlying member method code with example usage: obj.testTopK(dtype, rank) and returns: nn_ops