top_k

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

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

axis could be negative:

# A negative axis is the last-k axis
x = dragon.constant([[1, 2, 3], [3, 2, 1]])
value1, index1 = dragon.math.top_k(x, k=2, axis=1)
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])
_, index1 = dragon.math.top_k(-x, largest=True)
_, index2 = dragon.math.top_k(x, largest=False)  # Equivalent
Parameters:
  • inputs (dragon.Tensor) – The input tensor.
  • k (int, optional, default=1) – The number of top elements to select.
  • axis (int, optional, default=-1) – The axis to retrieve.
  • largest (bool, optional, default=True) – Return largest or smallest elements.
  • sorted (bool, optional, default=True) – Whether to return in the sorted order.
Returns:

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