torch.argsort
function argsort<S extends Shape, D extends DType = DType, Dev extends DeviceType = DeviceType>(input: Tensor<S, D, Dev>): Tensor<S, 'int32', Dev>function argsort<S extends Shape, D extends DType = DType, Dev extends DeviceType = DeviceType>(input: Tensor<S, D, Dev>, dim: number, descending: boolean, options?: SortOptions): Tensor<S, 'int32', Dev>Returns the indices that would sort a tensor along a given dimension.
Like sort but returns only indices, not values. Useful when you only need the ranking/ordering and don't need the actual sorted values. More efficient than sort when values aren't needed.
- Memory efficient: Returns only indices; much cheaper than sort when values not needed
- Discrete output: No gradients through indices (indices are discrete)
- Dimension semantics: -1 is last dimension; Python-style negative indexing supported
- Indices are positions: Element at x[indices[i]] gives i-th sorted value
- CPU/GPU: Runs on same device as input; optimized for both
Parameters
inputTensor<S, D, Dev>- Tensor to get sort indices from
Returns
Tensor<S, 'int32', Dev>– Indices tensor of same shape as input, showing sort orderExamples
// Get ranking indices
const x = torch.tensor([3, 1, 4, 1, 5, 9]);
const idx = torch.argsort(x);
// [1, 3, 0, 2, 4, 5] - elements at these positions give sorted order
// Descending order indices
const idx_desc = torch.argsort(x, { descending: true });
// [5, 4, 2, 0, 1, 3] - indices for largest to smallest
// Get top-k without topk function
const scores = torch.tensor([10, 20, 15, 5, 25]);
const top_idx = torch.argsort(scores, { descending: true }).slice(0, null, 3);
// First 3 indices of largest scores
// Ranking/permutation use case
const test_scores = torch.tensor([85, 92, 78, 92, 88]);
const rank_indices = torch.argsort(test_scores, { descending: true });
// Tell you who got 1st place (indices[0]), 2nd place (indices[1]), etc.
// Efficient batched ranking
const batch_scores = torch.randn(32, 100); // 32 samples, 100 features each
const rank_indices = torch.argsort(batch_scores, { dim: 1, descending: true });
// For each sample, indices showing feature ranking
// Reconstruction: get actual sorted values
const x = torch.tensor([3, 1, 4, 1, 5, 9]);
const idx = torch.argsort(x);
// To get sorted values: use gather operation with idxSee Also
- PyTorch torch.argsort()
- sort - Get both sorted values and indices
- topk - Get top k indices efficiently (faster for small k)
- argmax - Get index of maximum value only
- argmin - Get index of minimum value only