Introduction

In TensorFlow, argsort means returning the indices that would sort a tensor, not the sorted values themselves. The main API is tf.argsort, and it works across one-dimensional and multi-dimensional tensors. Once you have the indices, you can use them for ranking, reordering, top-element analysis, or aligning another tensor to the same sorted order.

Basic Usage

The simplest example is a one-dimensional tensor.

1import tensorflow as tf
2
3values = tf.constant([3, 1, 2])
4indices = tf.argsort(values)
5print(indices)

This returns indices equivalent to [1, 2, 0], because values[1] is the smallest element, then values[2], then values[0].

If you want the actual sorted values too, gather them using those indices.

sorted_values = tf.gather(values, indices)
print(sorted_values)

Ascending and Descending

By default, TensorFlow sorts in ascending order. You can reverse the direction explicitly.

1import tensorflow as tf
2
3values = tf.constant([3, 1, 2])
4print(tf.argsort(values, direction='ASCENDING'))
5print(tf.argsort(values, direction='DESCENDING'))

This is useful for ranking predictions or choosing the largest scores first.

Sorting Along an Axis

For matrices and higher-rank tensors, axis controls which dimension is sorted.

1import tensorflow as tf
2
3x = tf.constant([[3, 1, 2],
4                 [9, 7, 8]])
5
6print(tf.argsort(x, axis=-1))
7print(tf.argsort(x, axis=0))

• 'axis=-1 sorts within each row'
• 'axis=0 sorts down each column'

Reading the axis correctly matters more than the function call itself.

Stable Sorting

If equal values appear and their original order matters, use stable=True.

1import tensorflow as tf
2
3x = tf.constant([5, 2, 2, 9])
4print(tf.argsort(x, stable=True))

A stable sort preserves the original relative order of equal elements. That can matter when argsort is one step inside a larger ranking pipeline.

Common Pattern: Top Scores

For recommendation or classification tasks, you often argsort scores in descending order and take the first few indices.

1import tensorflow as tf
2
3scores = tf.constant([0.15, 0.92, 0.33, 0.81])
4ranked = tf.argsort(scores, direction='DESCENDING')
5top2 = ranked[:2]
6print(top2)

TensorFlow also has tf.math.top_k for the special case where you only need the top k items rather than a full ranking. That is usually a better choice when performance matters and full argsort is unnecessary.

Reordering Other Tensors with the Indices

One of the most useful things about argsort is that the returned indices can be applied to related tensors. For example, if one tensor stores scores and another stores labels, you can rank the labels by the score order.

1import tensorflow as tf
2
3scores = tf.constant([0.4, 0.9, 0.1])
4labels = tf.constant(['medium', 'high', 'low'])
5
6order = tf.argsort(scores, direction='DESCENDING')
7ranked_labels = tf.gather(labels, order)
8print(ranked_labels)

That pattern is common in ranking pipelines, recommendation systems, and evaluation code.

Argsort Versus Sort

It is easy to mix up tf.sort and tf.argsort.

• 'tf.sort returns sorted values'
• 'tf.argsort returns the indices that would produce that order'

1import tensorflow as tf
2
3values = tf.constant([3, 1, 2])
4print(tf.sort(values))
5print(tf.argsort(values))

If you need to reorder another tensor consistently with the sorted order of a score tensor, argsort is usually the right tool.

Common Pitfalls

• Expecting tf.argsort to return sorted values instead of sorted indices.
• Forgetting to set direction='DESCENDING' when ranking largest scores first.
• Sorting along the wrong axis in multi-dimensional tensors.
• Using full argsort when top_k would be cheaper and more direct.
• Ignoring sort stability when equal values must preserve their original order.

Summary

• 'tf.argsort returns indices, not values.'
• Use direction and axis deliberately to match the ranking you want.
• Combine tf.argsort with tf.gather when you also need the sorted values.
• Use stable=True if equal values should keep their original order.
• Reach for top_k when you only need the highest-ranked items rather than a full sort.