topk

dragon.vm.torch.topk(
  input,
  k,
  dim=None,
  largest=True,
  sorted=True,
  out=None
)[source]

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

If dim is not given, the last dimension is chosen:

x = torch.tensor([[1, 2, 3], [3, 2, 1]])
value1, index1 = torch.topk(x, k=2)
value2, index2 = torch.topk(x, k=2, dim=1)  # Equivalent

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

x = torch.tensor([[1, 2, 3], [3, 2, 1]])
_, index1 = torch.topk(x, 1, largest=False)
_, index2 = torch.topk(-x, 1, largest=True)  # Equivalent
Parameters:
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • k (int) – The number of top elements to select.
  • dim (int, optional) – The dimension to reduce.
  • largest (bool, optional) – Return largest or smallest elements.
  • sorted (bool, optional) – Whether to return in the sorted order.
  • out (Sequence[dragon.vm.torch.Tensor], optional) – The optional output value and index.
Returns:

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