NumPy argmax: np.argmax() Syntax, axis, keepdims and NaN

numpy.argmax returns the position of the largest value in an array, not the value itself. With no axis it counts positions in the flattened array; with an axis it returns one index for every slice along that axis. This guide covers the current signature, how the output shape follows from axis and keepdims, how to turn indices back into values or row-and-column coordinates, and the NaN trap that returns a wrong-looking answer without raising an error. Everything below follows the NumPy reference documentation, checked on 2026-10-29.

What np.argmax is used for in research

Researchers reach for np.argmax whenever the question is “where is the maximum?” rather than “what is the maximum?”: the time point of a peak in a signal, the grid cell with the highest score, or the most probable class in each row of a probability array. A worked example on this site, the categorical distribution guide, uses draws.argmax(axis=-1) to recover category labels from one-hot multinomial draws.

The function is part of NumPy’s array toolkit. If you are new to NumPy arrays, the NumPy array guide covers creating and shaping them first.

Install or access

NumPy’s own install page (numpy.org/install) lists two routes.

• pip: pip install numpy
• conda: conda create -n my-env, then conda activate my-env, then conda install numpy

To keep a pip install in an isolated environment, the same page shows python -m venv my-env, then source my-env/bin/activate on macOS and Linux or my-env\Scripts\activate on Windows, then pip install numpy.

In code, NumPy is imported by convention as import numpy as np; the NumPy beginners’ guide notes that this widespread convention gives access to NumPy features through the short np. prefix.

Syntax of numpy.argmax

numpy.argmax(a, axis=None, out=None, *, keepdims=<no value>)

• a: array_like, the input array.
• axis: int, optional. By default the index is into the flattened array; otherwise it is along the specified axis.
• out: array, optional. If provided, the result is inserted into this array, which should have the appropriate shape and dtype.
• keepdims: bool, optional. Reduced axes are left in the result as size-one dimensions so the result broadcasts against the input array. New in NumPy 1.22.0, and keyword-only (it follows the * in the signature). Older versioned documentation pages that still appear in search results predate this argument.

The return value, index_array, is an ndarray of ints with the same shape as a.shape with the dimension along axis removed. With keepdims=True the size of that axis is 1 instead.

Step-by-step: the five calls that cover most uses

The examples below are the ones given in the NumPy documentation for numpy.argmax.

1. Build an array.

   import numpy as np
   a = np.arange(6).reshape(2, 3) + 10
   # [[10, 11, 12],
   #  [13, 14, 15]]

2. Index into the flattened array (no axis).

   np.argmax(a)        # 5

   The maximum, 15, is the sixth element when the array is read row by row, so its flat index is 5.

3. Reduce along an axis.

   np.argmax(a, axis=0)   # array([1, 1, 1])
   np.argmax(a, axis=1)   # array([2, 2])

   axis=0 compares values down the rows and gives one index per column (three results for this 2 by 3 array). axis=1 compares values across the columns and gives one index per row (two results).

4. Handle ties.

   b = np.arange(6)
   b[1] = 5               # [0, 5, 2, 3, 4, 5]
   np.argmax(b)           # 1

   When the maximum occurs more than once, the indices of the first occurrence are returned. Only the first occurrence is returned.

5. Keep the reduced axis with keepdims.

   x = np.arange(24).reshape((2, 3, 4))
   res = np.argmax(x, axis=1, keepdims=True)
   res.shape              # (2, 1, 4)

Output and how to read it

The result always has the input’s shape minus the axis you reduced, unless you ask for keepdims=True.

Input shape	Call	Result
(2, 3)	np.argmax(a)	5, an index into the flattened array
(2, 3)	np.argmax(a, axis=0)	array([1, 1, 1]), shape (3,)
(2, 3)	np.argmax(a, axis=1)	array([2, 2]), shape (2,)
(2, 3, 4)	np.argmax(x, axis=1, keepdims=True)	shape (2, 1, 4)

Get the value instead of the index

For the flattened case, use the flat iterator: a.flat[np.argmax(a)] returns 15 for the array above. Along an axis, the documentation pairs np.argmax with np.take_along_axis:

x = np.array([[4, 2, 3], [1, 0, 3]])
index_array = np.argmax(x, axis=-1)
np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)
# array([[4],
#        [3]])

The documentation notes this is the same as np.amax(x, axis=-1, keepdims=True). Adding .squeeze(axis=-1) to the take_along_axis call gives array([4, 3]).

Get the row and column of the maximum

Convert the flat index to coordinates with np.unravel_index:

ind = np.unravel_index(np.argmax(a, axis=None), a.shape)
ind        # (1, 2)
a[ind]     # 15

Troubleshooting

The result is the position of a NaN

np.argmax does not skip missing values. The NumPy nanargmax documentation shows the difference on an array with one NaN:

a = np.array([[np.nan, 4], [2, 3]])
np.argmax(a)       # 0, the position of the NaN
np.nanargmax(a)    # 1

If your measurements can contain missing values, check for them first or use np.nanargmax. For detecting and replacing NaN, see NumPy NaN: how to detect, ignore and replace np.nan.

nanargmax raises ValueError on an all-NaN slice

np.nanargmax returns the indices of the maximum values along the specified axis, ignoring NaNs, and for all-NaN slices a ValueError is raised. The documentation also warns that the results cannot be trusted if a slice contains only NaNs and -Infs.

Getting one index per column when you expected one per row

The axis you pass is the one that disappears. For a 2 by 3 array, axis=0 returns three indices (one for each column) and axis=1 returns two (one for each row), as in the table above.

Mistaking an index for a value

np.argmax returns where the maximum is. To get the maximum itself use the patterns in “Get the value instead of the index”, or np.amax.

keepdims is not in the documentation you are reading

The keepdims argument was added in NumPy 1.22.0. If you are reading an older versioned documentation page, it will not list it; use the stable documentation linked in the sources below.

Free alternatives and related functions

All of these ship with NumPy, which is free software; each row below comes from the NumPy documentation cited in this guide.

Function or pattern	What it does	Use it when
np.nanargmax(a, axis=None, out=None, *, keepdims=<no value>)	Index of the maximum, ignoring NaNs; raises ValueError for all-NaN slices	Data can contain missing values
a.flat[np.argmax(a)]	Value at the flat index	You need the maximum value from a flattened argmax
np.unravel_index(np.argmax(a), a.shape)	Row and column (N-dimensional coordinates) of the maximum	You need coordinates, not a flat index
np.take_along_axis(x, np.expand_dims(index_array, axis=-1), axis=-1)	Values at the argmax indices along an axis	You need the per-row maxima after an axis-wise argmax
np.amax(x, axis=-1, keepdims=True)	The maximum values themselves (same result as the take_along_axis pattern in the NumPy example)	You only need the values

For sorted order rather than just the top position, see NumPy argsort, and for the distinct values in an array see NumPy unique.

FAQ

What does np.argmax return?

The index of the maximum value, as an integer array. Without axis it is an index into the flattened array; with axis it has the input’s shape with that axis removed.

What if the maximum appears more than once?

The indices of the first occurrence are returned. In the documentation example [0, 5, 2, 3, 4, 5], np.argmax returns 1.

How do I get the value instead of the index?

Use a.flat[np.argmax(a)] for the whole array, or np.take_along_axis with np.expand_dims along an axis; np.amax returns the values directly.

How do I get the row and column of the maximum?

Pass the flat index to np.unravel_index with the array’s shape: np.unravel_index(np.argmax(a, axis=None), a.shape) returns (1, 2) for the example array.

How do I ignore NaN values?

Use np.nanargmax. On [[nan, 4], [2, 3]] it returns 1, where np.argmax returns 0.

What does keepdims do?

It leaves the reduced axis in the result as a size-one dimension so the result broadcasts against the input. For shape (2, 3, 4) with axis=1 the result shape is (2, 1, 4).

Sources

• NumPy documentation, numpy.argmax, checked 2026-10-29.
• NumPy documentation, numpy.nanargmax, checked 2026-10-29.
• NumPy, Installing NumPy, checked 2026-10-29.
• NumPy documentation, NumPy: the absolute basics for beginners, checked 2026-10-29.