top_k

dragon.math.top_k(
  inputs,
  k=1,
  axis=None,
  largest=True,
  sorted=True,
  **kwargs
)[source]

Return the top-K largest or smallest elements along the given axis.

If axis is not given, the last axis is chosen:

x = dragon.constant([[1, 2, 3], [3, 2, 1]])
value1, index1 = dragon.math.top_k(x, k=2)
value2, index2 = dragon.math.top_k(x, k=2, axis=1)  # Equivalent

If largest is False, the k smallest elements are returned:

x = dragon.constant([[1, 2, 3], [3, 2, 1]])
_, index1 = dragon.math.top_k(x, largest=False)
_, index2 = dragon.math.top_k(-x, largest=True)  # Equivalent
Parameters:
  • inputs (dragon.Tensor) – The input tensor.
  • k (int, optional, default=1) – The number of top elements to select.
  • axis (int, optional) – The axis to reduce.
  • largest (bool, optional, default=True) – Return largest or smallest elements.
  • sorted (bool, optional) – Whether to return in the sorted order.
Returns:

Sequence[dragon.vm.torch.Tensor] – The value and index tensor.