argmax

argmax(x: array, /, *, axis: int | None = None, keepdims: bool = False) → array

Returns the indices of the maximum values along a specified axis.

When the maximum value occurs multiple times, only the indices corresponding to the first occurrence are returned.

Parameters:

• x (array) – input array. Should have a real-valued data type.

• axis (Optional[int]) – axis along which to search. If None, the function must return the index of the maximum value of the flattened array. If not None, a valid axis must be an integer on the interval [-N, N), where N is the number of axes in x. If an axis is specified as a negative integer, the function must determine the axis along which to perform the operation by counting backward from the last axis (where -1 refers to the last axis). If provided an invalid axis, the function must raise an exception. Default: None.

• keepdims (bool) – if True, the reduced axes must be included in the result as singleton dimensions, and, accordingly, the result must be broadcast-compatible with the input array (see Broadcasting). Otherwise, if False, the reduced axes must not be included in the result. Default: False.

Returns:

out (array) – if axis is None, a zero-dimensional array containing the index of the first occurrence of the maximum value; otherwise, a non-zero-dimensional array containing the indices of the maximum values. The returned array must have be the default array index data type.

Notes

• For backward compatibility, conforming implementations may support complex numbers; however, inequality comparison of complex numbers is unspecified and thus implementation-dependent (see Complex Number Ordering).