def _batch_sort_vector(x, ascending=True, name=None):
  with ops.name_scope(name, "sort_each_row", [x]):
    x = ops.convert_to_tensor(x, name="x")
    n = array_ops.shape(x)[-1]
    if ascending:
      y, _ = nn_ops.top_k(-x, k=n, sorted=True)
      y = -y
    else:
      y, _ = nn_ops.top_k(x, k=n, sorted=True)
    y.set_shape(x.shape)
    return y

  def get_best(self, n):
    """Return the indices and values of the n highest scores in the TopN."""

    def refresh_shortlist():
      """Update the shortlist with the highest scores in id_to_score."""
      new_scores, new_ids = nn_ops.top_k(self.id_to_score, self.shortlist_size)
      smallest_new_score = math_ops.reduce_min(new_scores)
      new_length = math_ops.reduce_sum(
          math_ops.to_int32(math_ops.greater(new_scores, dtypes.float32.min)))
      u1 = self.sl_ids.assign(
          math_ops.to_int64(array_ops.concat([[new_length], new_ids], 0)))
      u2 = self.sl_scores.assign(
          array_ops.concat([[smallest_new_score], new_scores], 0))
      self.last_ops = [u1, u2]
      return control_flow_ops.group(u1, u2)

    # We only need to refresh the shortlist if n is greater than the
    # current shortlist size (which is stored in sl_ids[0]).
    with ops.control_dependencies(self.last_ops):
      cond_op = control_flow_ops.cond(n > self.sl_ids[0], refresh_shortlist,
                                      control_flow_ops.no_op)
      with ops.control_dependencies([cond_op]):
        topk_values, topk_indices = nn_ops.top_k(
            self.sl_scores,
            math_ops.minimum(n, math_ops.to_int32(self.sl_ids[0])))
        # topk_indices are the indices into the shortlist, we want to return
        # the indices into id_to_score
        gathered_indices = array_ops.gather(self.sl_ids, topk_indices)
        return gathered_indices, topk_values

 def testKNegative(self):
   inputs = [[0.1, 0.2], [0.3, 0.4]]
   with self.test_session(use_gpu=True):
     k = array_ops.placeholder(dtypes.int32)
     values, _ = nn_ops.top_k(inputs, k)
     with self.assertRaisesOpError("Need k >= 0, got -7"):
       values.eval(feed_dict={k: -7})

 def testTopKGradients(self):
   with self.test_session(use_gpu=True) as sess:
     inputs = array_ops.placeholder(dtypes.int32, shape=[2, 5])
     values, _ = nn_ops.top_k(inputs, 3)
     grad = sess.run(
         gradients_impl.gradients(
             values, inputs, grad_ys=[[[1, 2, 3], [4, 5, 6]]]),
         feed_dict={inputs: [[2, -1, 1000, 3, 4], [1, 5, 2, 4, 3]]})[0]
   self.assertEqual(grad.tolist(), [[0, 0, 1, 3, 2], [0, 4, 0, 5, 6]])

 def testTopKGradients(self):
   with self.session(use_gpu=True) as sess:
     inputs = array_ops.placeholder(dtypes.float32, shape=[2, 5])
     values, _ = nn_ops.top_k(inputs, 3)
     grad = sess.run(
         gradients_impl.gradients(
             values, inputs, grad_ys=[[[1., 2., 3.], [4., 5., 6.]]]),
         feed_dict={inputs: [[2., -1., 1000., 3., 4.],
                             [1., 5., 2., 4., 3.]]})[0]
   self.assertEqual(
       grad.tolist(), [[0., 0., 1., 3., 2.], [0., 4., 0., 5., 6.]])

 def refresh_shortlist():
   """Update the shortlist with the highest scores in id_to_score."""
   new_scores, new_ids = nn_ops.top_k(self.id_to_score, self.shortlist_size)
   smallest_new_score = math_ops.reduce_min(new_scores)
   new_length = math_ops.reduce_sum(
       math_ops.to_int32(math_ops.greater(new_scores, dtypes.float32.min)))
   u1 = self.sl_ids.assign(
       math_ops.to_int64(array_ops.concat([[new_length], new_ids], 0)))
   u2 = self.sl_scores.assign(
       array_ops.concat([[smallest_new_score], new_scores], 0))
   self.last_ops = [u1, u2]
   return control_flow_ops.group(u1, u2)

def _sort_rows(matrix, num_rows):
  """Sort matrix rows by the last column.

  Args:
      matrix: a matrix of values (row,col).
      num_rows: (int) number of sorted rows to return from the matrix.
  Returns:
      Tensor (num_rows, col) of the sorted matrix top K rows.
  """
  tmatrix = array_ops.transpose(matrix, [1, 0])
  sorted_tmatrix = nn_ops.top_k(tmatrix, num_rows)[0]
  return array_ops.transpose(sorted_tmatrix, [1, 0])

  def _validateTopK(self,
                    inputs,
                    k,
                    expected_values,
                    expected_indices,
                    sorted=True):  # pylint: disable=redefined-builtin
    np_expected_values = np.array(expected_values)
    np_expected_indices = np.array(expected_indices)
    with self.test_session(use_gpu=True) as sess:
      values_op, indices_op = nn_ops.top_k(inputs, k, sorted=sorted)
      values, indices = sess.run([values_op, indices_op])

      self.assertShapeEqual(np_expected_values, values_op)
      self.assertShapeEqual(np_expected_indices, indices_op)

      if sorted:
        self.assertAllClose(np_expected_values, values)
        # Do some special casing of equality of indices: if indices
        # are not the same, but values are floating type, ensure that
        # the values are within epsilon of each other.
        if not np.issubdtype(np_expected_values.dtype, np.float):
          # Values are not floating point type; check indices exactly
          self.assertAllEqual(np_expected_indices, indices)
        else:
          # Values are floating point; indices may be swapped for
          # values near each other.
          indices_not_equal = np_expected_indices != indices
          if np.any(indices_not_equal):
            values_unsure = values[indices_not_equal]
            expected_values_unsure = expected_values[indices_not_equal]
            self.assertAllClose(expected_values_unsure, values_unsure)
      else:
        np_inputs = np.array(inputs)

        # Check that the indices are valid.
        for result_index, src_index in np.ndenumerate(indices):
          value = values[result_index]
          expected_value = np_inputs[result_index[0], src_index]
          np.testing.utils.assert_almost_equal(value, expected_value)

        # Check that if two elements are equal, the lower-index element appears
        # first.
        shape = values.shape
        for batch_index in range(shape[0]):
          for index in range(shape[1] - 1):
            if np.isclose(values[batch_index, index],
                          values[batch_index, index + 1]):
              self.assertLess(indices[batch_index, index],
                              indices[batch_index, index + 1])

        # Now check the results, ignoring order.
        self.assertAllEqual(np.sort(np_expected_indices), np.sort(indices))
        self.assertAllClose(np.sort(np_expected_values), np.sort(values))

  def _update_mask(self, weights, threshold):
    """Updates the mask for a given weight tensor.

    This functions first computes the cdf of the weight tensor, and estimates
    the threshold value such that 'desired_sparsity' fraction of weights
    have magnitude less than the threshold.

    Args:
      weights: The weight tensor that needs to be masked.
      threshold: The current threshold value. The function will compute a new
        threshold and return the exponential moving average using the current
        value of threshold

    Returns:
      new_threshold: The new value of the threshold based on weights, and
        sparsity at the current global_step
      new_mask: A numpy array of the same size and shape as weights containing
        0 or 1 to indicate which of the values in weights falls below
        the threshold

    Raises:
      ValueError: if sparsity is not defined
    """
    if self._sparsity is None:
      raise ValueError('Sparsity variable undefined')

    sparsity = self._get_sparsity(weights.op.name)
    with ops.name_scope(weights.op.name + '_pruning_ops'):
      abs_weights = math_ops.abs(weights)
      k = math_ops.cast(
          math_ops.round(
              math_ops.cast(array_ops.size(abs_weights), dtypes.float32) *
              (1 - sparsity)), dtypes.int32)
      # Sort the entire array
      values, _ = nn_ops.top_k(
          array_ops.reshape(abs_weights, [-1]), k=array_ops.size(abs_weights))
      # Grab the (k-1) th value
      current_threshold = array_ops.gather(values, k - 1)
      smoothed_threshold = math_ops.add_n([
          math_ops.multiply(current_threshold, 1 - self._spec.threshold_decay),
          math_ops.multiply(threshold, self._spec.threshold_decay)
      ])

      new_mask = math_ops.cast(
          math_ops.greater_equal(abs_weights, smoothed_threshold),
          dtypes.float32)

    return smoothed_threshold, new_mask

 def _validateTopK(self,
                   inputs,
                   k,
                   expected_values,
                   expected_indices,
                   sorted=True):
   np_values = np.array(expected_values)
   np_indices = np.array(expected_indices)
   with self.test_session():
     values_op, indices_op = nn_ops.top_k(inputs, k, sorted=sorted)
     values = values_op.eval()
     indices = indices_op.eval()
     self.assertAllClose(np_values, values)
     self.assertAllEqual(np_indices, indices)
     self.assertShapeEqual(np_values, values_op)
     self.assertShapeEqual(np_indices, indices_op)

  def testTopKZeros(self):
    """Tests that positive and negative zeros sort correctly."""
    # Only bfloat16 is implemented.
    bfloat16 = dtypes.bfloat16.as_numpy_dtype
    if bfloat16 not in self.numeric_types:
      return

    with self.cached_session() as sess:
      p = array_ops.placeholder(dtypes.bfloat16)
      with self.test_scope():
        topk = nn_ops.top_k(p, k=4)
      results = sess.run(
          topk,
          {p: np.array([0., -0., 0., 3., -0., -4., 0., -0.], dtype=bfloat16)})
      self.assertAllEqual(
          np.array([3., 0., 0., 0.], dtype=bfloat16), results[0])
      self.assertEqual(list([3, 0, 2, 6]), list(results[1]))

def _filter_top_k(x, k):
  """Filters top-k values in the last dim of x and set the rest to NEG_INF.

  Used for computing top-k prediction values in dense labels (which has the same
  shape as predictions) for recall and precision top-k metrics.

  Args:
    x: tensor with any dimensions.
    k: the number of values to keep.

  Returns:
    tensor with same shape and dtype as x.
  """
  _, top_k_idx = nn_ops.top_k(x, k, sorted=False)
  top_k_mask = math_ops.reduce_sum(
      array_ops.one_hot(top_k_idx, x.shape[-1], axis=-1), axis=-2)
  return x * top_k_mask + NEG_INF * (1 - top_k_mask)

  def __init__(self, permutation, validate_args=False, name=None):
    """Creates the `Permute` bijector.

    Args:
      permutation: An `int`-like vector-shaped `Tensor` representing the
        permutation to apply to the rightmost dimension of the transformed
        `Tensor`.
      validate_args: Python `bool` indicating whether arguments should be
        checked for correctness.
      name: Python `str`, name given to ops managed by this object.

    Raises:
      TypeError: if `not permutation.dtype.is_integer`.
      ValueError: if `permutation` does not contain exactly one of each of
        `{0, 1, ..., d}`.
    """
    with ops.name_scope(name, "permute", values=[permutation]):
      permutation = ops.convert_to_tensor(
          permutation,
          name="permutation")
      if not permutation.dtype.is_integer:
        raise TypeError("permutation.dtype ({}) should be `int`-like.".format(
            permutation.dtype.name))
      p = tensor_util.constant_value(permutation)
      if p is not None:
        if set(p) != set(np.arange(p.size)):
          raise ValueError("Permutation over `d` must contain exactly one of "
                           "each of `{0, 1, ..., d}`.")
      elif validate_args:
        p, _ = nn_ops.top_k(-permutation,
                            k=array_ops.shape(permutation)[-1],
                            sorted=True)
        permutation = control_flow_ops.with_dependencies([
            check_ops.assert_equal(
                -p, math_ops.range(array_ops.size(p)),
                message=("Permutation over `d` must contain exactly one of "
                         "each of `{0, 1, ..., d}`.")),
        ], permutation)
      self._permutation = permutation
      super(Permute, self).__init__(
          forward_min_event_ndims=1,
          is_constant_jacobian=True,
          validate_args=validate_args,
          name=name or "permute")

 def GetParams(self):
   """Testing Top-K in TF-TRT conversion."""
   dtype = dtypes.float32
   input_name = "input"
   input_dims = [100, 100]
   k = 5
   g = ops.Graph()
   with g.as_default():
     x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name)
     k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
     values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
     values = array_ops.identity(values, name="output_values")
     indices = array_ops.identity(indices, name="output_indices")
   return trt_test.TfTrtIntegrationTestParams(
       gdef=g.as_graph_def(),
       input_names=[input_name],
       input_dims=[[input_dims]],
       output_names=["output_values", "output_indices"],
       expected_output_dims=[[[100, k], [100, k]]])

  def testTopKInfinities(self):
    """Tests that positive and negative infinity sort correctly."""
    # Only bfloat16 is implemented.
    bfloat16 = dtypes.bfloat16.as_numpy_dtype
    if bfloat16 not in self.numeric_types:
      return

    with self.cached_session() as sess:
      p = array_ops.placeholder(dtypes.bfloat16)
      with self.test_scope():
        topk = nn_ops.top_k(p, k=6)
      results = sess.run(topk, {
          p: np.array(
              [1, 2, float("inf"), -float("inf"), -1, -2], dtype=bfloat16)
      })
      self.assertAllEqual(
          np.array(
              [float("inf"), 2.0, 1.0, -1.0, -2.0, -float("inf")],
              dtype=bfloat16), results[0])
      self.assertEqual(list([2, 1, 0, 4, 5, 3]), list(results[1]))

  def testTopKZeros(self):
    """Tests that positive and negative zeros sort correctly."""
    # TODO(b/26783907): The Sort HLO is not implemented on CPU or GPU.
    if self.device in ["XLA_CPU", "XLA_GPU"]:
      return

    # Only bfloat16 is implemented.
    bfloat16 = dtypes.bfloat16.as_numpy_dtype
    if bfloat16 not in self.numeric_types:
      return

    with self.test_session() as sess:
      p = array_ops.placeholder(dtypes.bfloat16)
      with self.test_scope():
        topk = nn_ops.top_k(p, k=4)
      results = sess.run(
          topk,
          {p: np.array([0., -0., 0., 3., -0., -4., 0., -0.], dtype=bfloat16)})
      self.assertAllEqual(
          np.array([3., 0., 0., 0.], dtype=bfloat16), results[0])
      self.assertEqual(list([3, 0, 1, 2]), list(results[1]))

 def GetParams(self):
   """Testing that output type of engine using Top-K is set correctly."""
   dtype = dtypes.float32
   input_name = "input"
   input_dims = [100, 100]
   k = 5
   g = ops.Graph()
   with g.as_default():
     x = array_ops.placeholder(dtype=dtype, shape=input_dims, name=input_name)
     k_tensor = constant_op.constant(k, dtype=dtypes.int32, name="Const")
     values, indices = nn_ops.top_k(x, k_tensor, name="TopK")
     # Reshape will act as a layer between the TopK output and the engine
     # output, requiring the output tensor of reshape to be set explicitly to
     # int32.
     indices = array_ops.reshape(indices, [100, 1, 5], name="Reshape")
     values = array_ops.identity(values, name="output_values")
     indices = array_ops.identity(indices, name="output_indices")
   return trt_test.TfTrtIntegrationTestParams(
       gdef=g.as_graph_def(),
       input_names=[input_name],
       input_dims=[[input_dims]],
       output_names=["output_values", "output_indices"],
       expected_output_dims=[[[100, k], [100, 1, k]]])

 def benchmarkTopK(self):
   for (m, n, p, use_gpu) in itertools.product(
       [128],
       [10, 100, 1000, 10000, 100000],
       [0.001, 0.01, 0.5, 0.99, 1.0],
       [False, True]):
     k = int(p * n)
     if k == 0:
       continue
     name = "m_%d_n_%d_k_%g_use_gpu_%s" % (m, n, k, use_gpu)
     device = "/%s:0" % ("gpu" if use_gpu else "cpu")
     with ops.Graph().as_default():
       with ops.device(device):
         x = random_ops.random_uniform((m, n))
         v = resource_variable_ops.ResourceVariable(x)
         op = nn_ops.top_k(v, k)
       with session.Session() as sess:
         v.initializer.run()
         r = self.run_op_benchmark(sess, op, min_iters=100, name=name)
         gb_processed_input = m * n / 1.0e9
         throughput = gb_processed_input / r["wall_time"]
         print("Benchmark: %s \t wall_time: %0.03g s \t "
               "Throughput: %0.03g GB/s" % (name, r["wall_time"], throughput))
         sys.stdout.flush()

def tftop_k(_):
  x = array_ops.placeholder(dtypes.int32, shape=[5], name='x')
  output = nn_ops.top_k(x, 2, name='values')
  array_ops.identity(output[1], name='indices')

def _descending_sort(values, axis, return_argsort=False):
  """Sorts values in reverse using `top_k`.

  Args:
    values: Tensor of numeric values.
    axis: Index of the axis which values should be sorted along.
    return_argsort: If False, return the sorted values. If True, return the
      indices that would sort the values.

  Returns:
    The sorted values.
  """
  k = array_ops.shape(values)[axis]
  rank = array_ops.rank(values)
  static_rank = values.shape.ndims
  # Fast path: sorting the last axis.
  if axis == -1 or axis + 1 == values.get_shape().ndims:
    top_k_input = values
    transposition = None
  else:
    # Otherwise, transpose the array. Swap axes `axis` and `rank - 1`.
    if axis < 0:
      # Calculate the actual axis index if counting from the end. Use the static
      # rank if available, or else make the axis back into a tensor.
      axis += static_rank or rank
    if static_rank is not None:
      # Prefer to calculate the transposition array in NumPy and make it a
      # constant.
      transposition = constant_op.constant(
          np.r_[
              # Axes up to axis are unchanged.
              np.arange(axis),
              # Swap axis and rank - 1.
              [static_rank - 1],
              # Axes in [axis + 1, rank - 1) are unchanged.
              np.arange(axis + 1, static_rank - 1),
              # Swap axis and rank - 1.
              [axis]],
          name='transposition')
    else:
      # Generate the transposition array from the tensors.
      transposition = array_ops.concat(
          [
              # Axes up to axis are unchanged.
              math_ops.range(axis),
              # Swap axis and rank - 1.
              [rank - 1],
              # Axes in [axis + 1, rank - 1) are unchanged.
              math_ops.range(axis + 1, rank - 1),
              # Swap axis and rank - 1.
              [axis]
          ],
          axis=0)
    top_k_input = array_ops.transpose(values, transposition)

  values, indices = nn_ops.top_k(top_k_input, k)
  return_value = indices if return_argsort else values
  if transposition is not None:
    # transposition contains a single cycle of length 2 (swapping 2 elements),
    # so it is an involution (it is its own inverse).
    return_value = array_ops.transpose(return_value, transposition)
  return return_value