code: def _top_k(
      input, k=1, sorted=True, name=None
  ):  # pylint: disable=unused-argument,redefined-builtin,missing-docstring
      # This currently ignores sorted=False. However, this should be safe since
      # call sites that don't invoke sorted=True will assume results are unsorted
      # (vs. never sorted), and so sorted results shouldn't impact results.
      input = _convert_to_tensor(input)
      if JAX_MODE:
          # JAX automatically returns values in sorted order.
          return TopK(*jax.lax.top_k(input, k))
      n = int(input.shape[-1] - 1)
      # For the values, we sort the negative entries and choose the smallest ones
      # and negate. This is equivalent to choosing the largest entries
      values = -np.sort(-input, axis=-1)[..., :k]
      # For the indices, we could argsort and reverse the entries and choose the
      # first k entries. However, this does not work in the case of ties, since the
      # first index a value occurs at is preferred. Thus we also reverse the input
      # to ensure the last tied value becomes first, and subtract this off from the
      # last index since the list is reversed.
      indices = (
          n
          - (
              np.argsort(input[..., ::-1], kind="stable", axis=-1)[
                  ..., ::-1
              ]
          )
      )[..., :k].astype(np.int32)
      return TopK(values, indices)
  
imports: import collections
  import functools
  import numpy
  import numpy
  from tensorflow_probability.python.internal.backend.numpy import _utils
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import _reverse
  from tensorflow_probability.python.internal.backend.numpy.numpy_array import one_hot
  from tensorflow_probability.python.internal.backend.numpy.ops import _convert_to_tensor
label_desc: Use image node_id 22 for calling a global function with example usage: _top_k(input, k, sorted, name) and returns: TopK, TopK