jax.numpy.argmin#

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

Return the index of the minimum value of an array.

JAX implementation ofnumpy.argmin().

Parameters:

• a (ArrayLike) – input array

• axis (int |None) – optional integer specifying the axis along which to find the minimumvalue. 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 minimum value along the specified axis.

Return type:

Array

Note

When the minimum value occurs more than once along a particular axis, thesmallest index is returned.

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,2])>>>jnp.argmin(x)Array(0, dtype=int32)

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

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