topk¶

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

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

dim could be negative:

# A negative dimension is the last-k dimension
x = torch.tensor([[1, 2, 3], [3, 2, 1]])
value1, index1 = torch.topk(x, k=2, dim=1)
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])
_, index1 = torch.topk(-x, 1, largest=True)
_, index2 = torch.topk(x, 1, largest=False)  # Equivalent

Parameters:

• input (dragon.vm.torch.Tensor) – The input tensor.
• k (int) – The number of top elements to select.
• dim (int, optional, default=-1) – The dimension to retrieve.
• largest (bool, optional) – Return largest or smallest elements.
• sorted (bool, optional) – Whether to return elements in the sorted order.
• out (Sequence[dragon.vm.torch.Tensor], optional) – The output value and index tensor.

Returns:

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