jax.numpy.nanargmax

jax.numpy.nanargmax(a, axis=None)

Return the indices of the maximum values in the specified axis ignoring

LAX-backend implementation of nanargmax().

Warning: jax.numpy.argmax returns -1 for all-NaN slices and does not raise an error.

Original docstring below.

NaNs. For all-NaN slices ValueError is raised. Warning: the results cannot be trusted if a slice contains only NaNs and -Infs.

Parameters

• a (array_like) – Input data.

• axis (int, optional) – Axis along which to operate. By default flattened input is used.

Returns

index_array – An array of indices or a single index value.

Return type

ndarray