jax.numpy.nanargmax#

jax.numpy.nanargmax(a,axis=None,out=None,keepdims=None)[source]#

Return the index of the maximum value of an array, ignoring NaNs.

JAX implementation ofnumpy.nanargmax().

Parameters:

• a (ArrayLike) – input array

• axis (int |None) – optional integer specifying the axis along which to find the maximumvalue. Ifaxis is not specified,a will be flattened.

• out (None) – unused by JAX

• keepdims (bool |None) – if True, then return an array with the same number of dimensionsasa.

Returns:

an array containing the index of the maximum value along the specified axis.

Return type:

Array

Note

In the case of an axis with all-NaN values, the returned index will be -1.This differs from the behavior ofnumpy.nanargmax(), which raises an error.

See also

• jax.numpy.argmax(): return the index of the maximum value.

• jax.numpy.nanargmin(): computeargmin while ignoring NaN values.

Examples

>>>x=jnp.array([1,3,5,4,jnp.nan])

Using a standardargmax() leads to potentially unexpected results:

>>>jnp.argmax(x)Array(4, dtype=int32)

Usingnanargmax returns the index of the maximum non-NaN value.

>>>jnp.nanargmax(x)Array(2, dtype=int32)

>>>x=jnp.array([[1,3,jnp.nan],...[5,4,jnp.nan]])>>>jnp.nanargmax(x,axis=1)Array([1, 0], dtype=int32)

>>>jnp.nanargmax(x,axis=1,keepdims=True)Array([[1],       [0]], dtype=int32)