Introduction

To get the index of a maximum value in TensorFlow, the main tool is tf.argmax. The part that matters most is not the function name but the axis you choose, because that determines whether you are asking for the maximum per row, per column, or across the whole tensor.

Use tf.argmax for One-Dimensional Tensors

For a one-dimensional tensor, tf.argmax returns the index of the largest value.

1import tensorflow as tf
2
3x = tf.constant([1.5, 7.2, 3.0, 7.1])
4idx = tf.argmax(x)
5
6print(idx.numpy())

This returns 1, because 7.2 is the largest value and it sits at index 1.

Axis Selection Is the Real Question

For multi-dimensional tensors, tf.argmax reduces along one axis and returns the index of the maximum inside each slice.

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

Those two calls mean different things:

• 'axis=0 compares values down each column'
• 'axis=1 compares values across each row'

Choosing the wrong axis is the most common reason people get a correct-looking but wrong answer.

Find the Global Maximum in a Matrix

Sometimes you want the one location of the maximum across all elements. In that case, flatten the tensor first and then convert the flat index back into coordinates.

1import tensorflow as tf
2
3x = tf.constant([
4    [1.0, 3.0, 2.0],
5    [4.0, 0.5, 6.0],
6], dtype=tf.float32)
7
8flat_index = tf.argmax(tf.reshape(x, [-1]))
9coords = tf.unravel_index(flat_index, tf.shape(x))
10
11print(flat_index.numpy())
12print(coords.numpy())

That gives both the flat offset and the row-column location of the overall maximum.

A Typical Classification Example

A very common use case is taking the predicted class from model output scores or logits.

1import tensorflow as tf
2
3logits = tf.constant([
4    [0.1, 1.9, 0.4],
5    [2.2, 0.3, 0.7],
6], dtype=tf.float32)
7
8predicted_classes = tf.argmax(logits, axis=1)
9print(predicted_classes.numpy())

If each row represents one example and each column one class, axis=1 returns the predicted class index for each example.

If you also want the winning score, combine argmax with reduce_max.

1probs = tf.nn.softmax(logits, axis=1)
2classes = tf.argmax(probs, axis=1)
3confidences = tf.reduce_max(probs, axis=1)
4
5print(classes.numpy())
6print(confidences.numpy())

Tie Behavior Matters

If more than one element shares the maximum value, TensorFlow returns the first index along the chosen axis.

1import tensorflow as tf
2
3x = tf.constant([4.0, 8.0, 8.0, 2.0])
4print(tf.argmax(x).numpy())

This returns 1, not 2, because the first maximum wins. That matters if downstream logic assumes ties are unique or random.

Use top_k When One Maximum Is Not Enough

If you need the top few indices rather than only one, use tf.math.top_k instead of repeated argmax calls.

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

That expresses the intent more clearly and avoids awkward repeated reduction logic.

Common Pitfalls

• Picking the wrong axis and then interpreting the result as if it meant something else.
• Forgetting that argmax returns the first maximum when values are tied.
• Treating the returned index tensor as if it were the maximum value itself.
• Forgetting to flatten first when the real goal is the global maximum over all elements.
• Using repeated argmax calls when the requirement is really top-k ranking.

Summary

• 'tf.argmax returns index positions, not the maximum values themselves.'
• The selected axis determines the meaning of the result.
• Flatten plus tf.unravel_index gives the global maximum location in multi-dimensional tensors.
• 'axis=1 is common for classification outputs where each row is one example.'
• Use tf.math.top_k when you need more than one best index.