fn

argsort

Tensor
argsort(input: Tensor, dim: DimLike = ...)
source

Return the indices that would sort input along dim.

Indexing input with the returned tensor along dim yields a sorted tensor. Equivalent to sort(input, dim)[1] but returns only the index half.

Parameters

inputTensor
Source tensor.
dimDimLike= ...
Axis along which to sort.

Returns

Tensor

int64 index tensor of the same shape as input.

Notes

The sort is stable: equal keys preserve their input order. For descending order, negate the input or pass descending=True to sort and use its index output.

Examples

>>> import lucid
>>> x = lucid.tensor([3, 1, 4, 1, 5])
>>> lucid.argsort(x, dim=0)
Tensor([1, 3, 0, 2, 4])