top_k¶

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

Return the top k-largest or k-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 elements in the sorted order.

Returns:

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