topk

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

Return the k largest/smallest values and indices along the given axis.

If dim is not given, the last dimension of the input is chosen.

If largest is False then the k smallest elements are returned.

Parameters:
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • k (int) – The top k.
  • dim (int, optional) – The axis of tensor to compute reduce value.
  • largest (bool, optional) – Whether to return largest or smallest elements.
  • sorted (bool, optional) – Whether to return in the sorted order.
  • out (dragon.vm.torch.Tensor, optional) – The optional output tensor.
Returns:

Sequence[dragon.vm.torch.Tensor] – The values and indices.