Introduction

Finding the top n largest or smallest values in a tensor is a common step in model post-processing, ranking, beam search, and feature analysis. In TensorFlow, the best tool depends on whether you need only the top results, all sorted indices, or control over the axis being processed.

Use tf.math.top_k for the Largest Values

If you need the indices of the largest n elements, tf.math.top_k is the most direct API. It returns both the values and the indices for the last dimension.

1import tensorflow as tf
2
3scores = tf.constant([0.3, 0.9, 0.2, 0.8, 0.7])
4values, indices = tf.math.top_k(scores, k=3)
5
6print(values.numpy())
7print(indices.numpy())

Typical output is:

[0.9 0.8 0.7]
[1 3 4]

That means the three largest values are at positions 1, 3, and 4.

tf.math.top_k is usually better than sorting the full tensor when you only need a small number of top results.

Get the Smallest Values with tf.argsort

For the smallest values, the most readable solution is often tf.argsort in ascending order, followed by slicing.

1import tensorflow as tf
2
3scores = tf.constant([0.3, 0.9, 0.2, 0.8, 0.7])
4indices = tf.argsort(scores, direction="ASCENDING")[:3]
5values = tf.gather(scores, indices)
6
7print(indices.numpy())
8print(values.numpy())

This returns the indices of the three smallest values, then gathers the values themselves.

You can also use tf.argsort(..., direction="DESCENDING")[:n] for the top maximum values if you want a single sorting-based pattern for both cases.

Using tf.math.top_k for Minimum Values

Another common trick is to negate the tensor. The smallest original values become the largest negated values.

1import tensorflow as tf
2
3scores = tf.constant([0.3, 0.9, 0.2, 0.8, 0.7])
4neg_values, indices = tf.math.top_k(-scores, k=2)
5min_values = -neg_values
6
7print(indices.numpy())
8print(min_values.numpy())

This works well when you want the efficiency of top_k but for minimum values.

Working with Matrices and Batches

Both tf.math.top_k and tf.argsort operate along an axis. tf.math.top_k uses the last dimension, which is convenient for batched model outputs.

1import tensorflow as tf
2
3logits = tf.constant([
4    [0.1, 0.8, 0.4, 0.7],
5    [0.9, 0.2, 0.5, 0.3],
6])
7
8values, indices = tf.math.top_k(logits, k=2)
9print(values.numpy())
10print(indices.numpy())

Each row gets its own top two entries. This is especially common in classification tasks where each row represents one example and each column represents one class.

If you need a different axis, tf.argsort is often more flexible because it accepts an axis parameter directly.

sorted_indices = tf.argsort(logits, axis=1, direction="DESCENDING")[:, :2]
print(sorted_indices.numpy())

Which Function Should You Choose?

Use tf.math.top_k when:

• You only need the largest n values
• You want both values and indices
• The last axis is the correct axis

Use tf.argsort when:

• You need full sorted order
• You want ascending and descending with one API
• You need to slice along a custom axis more explicitly

Both are valid; the right one depends on your shape and performance requirements.

Common Pitfalls

A frequent mistake is forgetting that tf.math.top_k works on the last dimension. If your ranking dimension is elsewhere, the output can look wrong even though the operation is technically valid.

Another mistake is slicing values but forgetting to gather the original tensor when using tf.argsort. tf.argsort returns indices, not the values themselves.

A third mistake is using full sorting when only the top few entries are needed. That can be slower than tf.math.top_k, especially for large tensors.

Summary

• Use tf.math.top_k to get the largest n values and indices efficiently.
• Use tf.argsort when you want sorted indices for minimum or maximum selection.
• For smallest values, either sort ascending or apply top_k to the negated tensor.
• Pay attention to the axis, especially for batched tensors.
• Gather values explicitly when your workflow starts from indices.