jax.Array.argmax — JAX documentation (original) (raw)

• API Reference
• jax.Array.argmax

jax.Array.argmax#

abstract Array.argmax(axis=None, out=None, keepdims=None)[source]#

Return the index of the maximum value.

Refer to jax.numpy.argmax() for the full documentation.

Parameters:

• self (Array)
• axis (int | None)
• out (None)
• keepdims (bool | None)

Return type:

Array