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