np.argmax() mirrors np.argmin() exactly, but finds the position of the largest value instead of the smallest. It's especially common in machine learning, where a model's output is often an array of scores or probabilities for different classes, and argmax() identifies which class index received the highest score, the predicted class. Like argmin(), it returns a flat index by default for multi-dimensional arrays, and returns the first index found on a tie.

1Understanding np.argmax()

💡

argmax() is the standard way to convert a model's array of per-class scores into a single predicted class index — it's one of the most common NumPy calls in a machine learning inference pipeline.

editor.html

import numpy as np

scores = np.array([0.1, 0.7, 0.2])
print(np.argmax(scores))

localhost:3000

2Practical Example

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

editor.html

import numpy as np

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

localhost:3000

3Best Practices

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

1. Use np.argmax() to convert an array of scores/probabilities into a predicted class index, rather than manually looping to find the largest value's position

2. Specify the axis argument explicitly for batched predictions, e.g. axis=1 to get one predicted class per row of a batch of samples

3. Remember argmax() returns the first index on a tie — be aware if your data could realistically have exact ties in the maximum value

⚠️

Tip: argmax() is the standard way to convert a model's array of per-class scores into a single predicted class index — it's one of the most common NumPy calls in a machine learning inference pipeline.

editor.html

import numpy as np

scores = np.array([0.1, 0.7, 0.2])
print(np.argmax(scores))

localhost:3000

Examples

Example 01Basic Usage

import numpy as np

scores = np.array([0.1, 0.7, 0.2])
print(np.argmax(scores))

Example 02Advanced Example

import numpy as np

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

Best Practices

Use np.argmax() to convert an array of scores/probabilities into a predicted class index, rather than manually looping to find the largest value's position
Specify the axis argument explicitly for batched predictions, e.g. axis=1 to get one predicted class per row of a batch of samples
Remember argmax() returns the first index on a tie — be aware if your data could realistically have exact ties in the maximum value

Interview Question

In a machine learning context, why is np.argmax() applied along axis=1 for a batch of predictions rather than with no axis argument at all?▼

Hint: Think about the shape of a batch of predictions — samples as rows, class scores as columns.

A batch of predictions is typically shaped as (number of samples, number of classes), with each row holding one sample's per-class scores. Calling argmax() with no axis argument would flatten the entire 2D array into one long sequence and return a single index for the overall largest score across the whole batch, losing the fact that you actually want one predicted class per sample. Specifying axis=1 tells NumPy to find the largest value within each row independently, collapsing the class dimension, and returning one predicted class index per sample, matching the batch's row-per-sample structure.

Exercises

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

View Solution

import numpy as np

scores = np.array([0.1, 0.7, 0.2])
print(np.argmax(scores))

Frequently Asked Questions

In a machine learning context, why is np.argmax() applied along axis=1 for a batch of predictions rather than with no axis argument at all?▼

Related Functions

np-argminnp-maxnp-where