tf.argmax() returns the index of the largest value along a specified axis of a tensor, commonly used to convert a model's per-class scores into a single predicted class.

Syntax

tf.argmax(input, axis=None)

Deep Dive Course

argmax() finds the position of the maximum value along the given axis, rather than the value itself — for a model's output layer producing one score per possible class, argmax() along the class axis identifies which class received the highest score, the model's actual predicted class. Without specifying an axis, it defaults to operating on the last axis, unlike NumPy's np.argmax(), which defaults to flattening the entire array first, which conveniently matches the common case of a tensor shaped as (batch_size, num_classes), where the class scores are the last dimension.

1Understanding tf.argmax()

💡

tf.argmax() defaults to operating along the last axis, unlike NumPy's np.argmax(), which defaults to flattening the whole array first — this default conveniently matches the common (batch_size, num_classes) shape of classification model outputs, letting you call it with no axis argument in that specific case.

editor.html

import tensorflow as tf

scores = tf.constant([0.1, 0.7, 0.2])
print(tf.argmax(scores))

localhost:3000

2Practical Example

Here is a real-world application of tf.argmax() showing how it is used in production TensorFlow code.

editor.html

import tensorflow as tf

batch_scores = tf.constant([[0.1, 0.7, 0.2], [0.6, 0.1, 0.3]])
predictions = tf.argmax(batch_scores, axis=1)
print(predictions)

localhost:3000

3Best Practices

Follow these guidelines when working with tf.argmax():

1. Use tf.argmax() to convert a model's per-class score/probability output into a single predicted class index, rather than manually scanning for the largest value

2. Rely on the default last-axis behavior for standard (batch_size, num_classes) shaped model outputs, but specify axis explicitly for anything with a different shape convention

3. Cast the result to a Python int, or compare it against integer labels with matching dtype, when checking predictions against ground-truth labels

⚠️

Tip: tf.argmax() defaults to operating along the last axis, unlike NumPy's np.argmax(), which defaults to flattening the whole array first — this default conveniently matches the common (batch_size, num_classes) shape of classification model outputs, letting you call it with no axis argument in that specific case.

editor.html

import tensorflow as tf

scores = tf.constant([0.1, 0.7, 0.2])
print(tf.argmax(scores))

localhost:3000

Examples

Example 01Basic Usage

import tensorflow as tf

scores = tf.constant([0.1, 0.7, 0.2])
print(tf.argmax(scores))

Example 02Advanced Example

import tensorflow as tf

batch_scores = tf.constant([[0.1, 0.7, 0.2], [0.6, 0.1, 0.3]])
predictions = tf.argmax(batch_scores, axis=1)
print(predictions)

Best Practices

Use tf.argmax() to convert a model's per-class score/probability output into a single predicted class index, rather than manually scanning for the largest value
Rely on the default last-axis behavior for standard (batch_size, num_classes) shaped model outputs, but specify axis explicitly for anything with a different shape convention
Cast the result to a Python int, or compare it against integer labels with matching dtype, when checking predictions against ground-truth labels

Interview Question

Why does tf.argmax() default to operating along the last axis, unlike NumPy's np.argmax(), which flattens the array by default?▼

Hint: Think about the typical shape of a classification model's output, and which axis actually represents the classes.

A classification model's output is conventionally shaped as (batch_size, num_classes), with the class scores specifically living along the last dimension — TensorFlow's default of operating on the last axis directly matches this extremely common convention, letting you call argmax() on a batch of predictions with no axis argument and immediately get one predicted class index per example. NumPy's argmax(), designed as a general-purpose array function without any particular assumption about what a tensor's dimensions represent, instead defaults to the more conservative, unambiguous choice of flattening the whole array first, requiring an explicit axis argument whenever you want per-row or per-column behavior instead.

Exercises

MediumPractice using tf.argmax() in a real scenario.

View Solution

import tensorflow as tf

scores = tf.constant([0.1, 0.7, 0.2])
print(tf.argmax(scores))

Frequently Asked Questions

Why does tf.argmax() default to operating along the last axis, unlike NumPy's np.argmax(), which flattens the array by default?▼

Related Functions

tf-reduce-summetrics-accuracynp-argmax