torch.Tensor.Tensor.argsort
Tensor.argsort(): Tensor<S, 'int32', Dev>Tensor.argsort(dimOrOptions?: number | SortOptions, descendingOrOptions?: boolean | SortOptions, options?: SortOptions): Tensor<S, 'int32', Dev>Returns the indices that would sort this tensor (without actually sorting values).
Computes the permutation that sorts the tensor along a dimension. Returns only the indices, not the sorted values. Use when you need the sorting order but want to apply it to related arrays (e.g., sorting loss by score index, reordering based on ranks).
Use Cases:
- Getting sort indices for reordering related data (e.g., sort images by confidence scores)
- Ranking elements (positions in sorted order)
- Implementing custom sorting logic with indices
- Reordering batches by loss, confidence, or other criteria
- Non-maximum suppression style operations (keep indices of top-k)
- Indices only: Returns indices, not sorted values (use sort() for both)
- Ascending default: Default sorts smallest-to-largest (use descending=true for largest-first)
- Dimension-specific: Only sorts along one dimension
- Stable: Equal elements maintain relative order
- 1-indexed not needed: Returns 0-indexed positions
- Batch support: Works independently on each batch element
- Index reuse: Indices can be reused to reorder other tensors
- Size tracking: Must track which elements correspond to which indices
Returns
Tensor<S, 'int32', Dev>– Indices tensor showing where each element would go in sorted orderExamples
// Get indices that sort values
const scores = torch.tensor([0.9, 0.3, 0.7, 0.5]);
const sorted_indices = scores.argsort(); // [1, 3, 2, 0] (ascending)
// Sort in descending order
const top_indices = scores.argsort(0, true); // [0, 2, 3, 1] (descending)
// Reorder related data using indices
const scores = torch.tensor([0.9, 0.3, 0.7, 0.5]);
const labels = torch.tensor([1, 2, 3, 4]);
const order = scores.argsort(0, true); // Sort by score descending
const sorted_labels = labels.index_select(0, order); // Reorder labels by score
// Rank elements
const values = torch.tensor([5, 2, 8, 1]);
const ranks = values.argsort().argsort(); // [2, 1, 3, 0] - rank of each element
// Sort along specific dimension in batches
const batch_scores = torch.randn(32, 10); // 32 samples, 10 scores each
const indices = batch_scores.argsort(1, true); // Get top-10 indices per batchSee Also
- PyTorch torch.argsort() (or tensor.argsort())
- sort - Get sorted values and indices together
- topk - Get top k values and indices (faster than full sort)
- kthvalue - Get k-th smallest value and its index
- index_select - Use indices to reorder elements