jax.numpy.argmax#
- jax.numpy.argmax(a, axis=None, out=None, keepdims=None)[source]#
Return the index of the maximum value of an array.
JAX implementation of
numpy.argmax()
.- Parameters:
a (ArrayLike) – input array
axis (int | None | None) – optional integer specifying the axis along which to find the maximum value. If
axis
is not specified,a
will be flattened.out (None | None) – unused by JAX
keepdims (bool | None | None) – if True, then return an array with the same number of dimensions as
a
.
- Returns:
an array containing the index of the maximum value along the specified axis.
- Return type:
See also
jax.numpy.argmin()
: return the index of the minimum value.jax.numpy.nanargmax()
: computeargmax
while ignoring NaN values.
Examples
>>> x = jnp.array([1, 3, 5, 4, 2]) >>> jnp.argmax(x) Array(2, dtype=int32)
>>> x = jnp.array([[1, 3, 2], ... [5, 4, 1]]) >>> jnp.argmax(x, axis=1) Array([1, 0], dtype=int32)
>>> jnp.argmax(x, axis=1, keepdims=True) Array([[1], [0]], dtype=int32)