jax.numpy.argmax ================ .. currentmodule:: jax.numpy .. autofunction:: argmax