Python np.argmax() function – AskPython

With `axis=0` on a 2×3 array, the result has one row index for each column. I like how that lets each column identify its own winning position.

Those indices can guide a separate lookup for each column’s winner. `np.argmax` returns positions within the values it searches, not the maximum values themselves.

What np.argmax returns

NumPy’s np.argmax() returns the zero-based index of a maximum value and, without an axis, searches the flattened array rather than returning a row-and-column coordinate. The result is an index, not the value.

import numpy as np
values = np.array([4, 9, 9])
print(np.argmax(values))

The output is 1 because the first 9 appears at index 1. Use np.max() or the array’s max method when you need the value instead of its position.

When two values tie, argmax returns the first matching position in the search order.

Read the result from the array shape

An axis tells NumPy which dimension to search across, and it determines which dimension disappears from the result.

The example uses a Python environment with NumPy installed. Run these commands from your project folder to create an isolated environment and install the package without pinning a version.

python3 -m venv .venv
.venv/bin/python -m pip install numpy

With axis 0, the result has one row index for each of the three columns. Axis 1 has one column index for each row.

For the input structure behind these examples, see NumPy arrays.

Choose a flat index, coordinates, or axis result

I ran the 2-by-3 array and got 4, not a coordinate pair. That number is a flat offset, so I passed it and the array shape to np.unravel_index().

I kept the input fixed across the example, so only the axis choice changes the row and column results.

import numpy as np

scores = np.array([[4, 2, 3], [1, 6, 2]])
flat_index = np.argmax(scores)
coordinates = np.unravel_index(flat_index, scores.shape)

print("flat index:", flat_index)
print("coordinates:", coordinates)
print("value at coordinates:", scores[coordinates])
print("axis=0:", np.argmax(scores, axis=0))
print("axis=1:", np.argmax(scores, axis=1))
print("axis=1 with keepdims:", np.argmax(scores, axis=1, keepdims=True))

out = np.empty(2, dtype=np.intp)
np.argmax(scores, axis=1, out=out)
print("axis=1 stored in out:", out)

tied = np.array([[7, 9, 9], [9, 2, 9]])
print("flat tie:", np.argmax(tied))
print("column ties:", np.argmax(tied, axis=0))

values = np.array([np.nan, 4.0])
print("argmax with NaN:", np.argmax(values))
print("nanargmax:", np.nanargmax(values))

try:
    np.argmax(np.array([]))
except ValueError as exc:
    print("empty input:", type(exc).__name__)

one_dimensional = np.array([4, 6, 6])
position = one_dimensional.argmax()
print("maximum and position:", one_dimensional[position], position)

index = np.argmax(scores)
winning_score = scores.flat[index]
print("maximum through flat index:", winning_score)

values = np.array([4, 9, 9])
print("first tied maximum index:", np.argmax(values))

Run the saved program from its folder to see each result in the terminal. The flat index, coordinates and axis outputs then appear together for comparison.

./.venv/bin/python argmax_demo.py

The flat result is 4, and coordinates (1, 1) point to its value, 6.

Axis 0 returns [0, 1, 0], the winning row within each column, while axis 1 returns [0, 1], the winning column within each row. Set keepdims=True to keep a length-one axis, which means the result has shape (2, 1) and can broadcast against the original input.

Supply an array through out when its shape and dtype match, and reserve np.unravel_index() for the flat result because axis calls already return local indices.

For another way to build an input, see how to reshape a NumPy array or create a sequence with arange.

Ties and special values

When equal values share the maximum, np.argmax returns the first one in the searched order. With an axis, “first” is evaluated inside each slice, so different rows or columns can select different positions.

• NaN values: NumPy 2.5.3 returned index 0 for [NaN, 4], while np.nanargmax returned 1. Use np.nanargmax() when NaNs should be ignored, and handle its ValueError for an all-NaN slice.
• Empty input: argmax cannot select a maximum from an empty array or an empty slice, so NumPy raises ValueError. Make sure every slice along the selected axis contains at least one value.

The NumPy documentation warns that nanargmax results cannot be trusted for a slice containing only NaNs and negative infinities. Read the nanargmax reference before relying on that boundary.

Use the index when position matters

Use argmax when the next operation needs the winning position, such as selecting a label or retrieving the value at that location. If the position does not matter, max returns the value directly, so there is no separate index to look up.

index = np.argmax(scores)
winning_score = scores.flat[index]

For the NumPy array in the example, this lookup returns 6. Keep the index when later code needs a coordinate. Otherwise, take the maximum value and leave the position out.

Frequently asked questions

The returned index refers to the values argmax searched, and the output shape changes with the axis choice.