argsort

dragon.vm.torch.argsort(
  input,
  dim=- 1,
  descending=False
)[source]

Return the index of sorted 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]])
index1 = torch.argsort(x, dim=1)
index2 = torch.argsort(x, dim=-1)  # Equivalent

Use descending to sort in the descending order:

x = torch.tensor([1, 2, 3])
index1 = torch.argsort(-x, descending=False)
index2 = torch.argsort(x, descending=True)  # Equivalent
Parameters:
  • input (dragon.vm.torch.Tensor) – The input tensor.
  • dim (int, optional, default=-1) – The dimension to sort.
  • descending (bool, optional, default=False) – Sort in the descending order or not.
Returns:

dragon.vm.torch.Tensor – The output tensor.