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
isFalse
, 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.