jax.lax.argmax
==============

.. currentmodule:: jax.lax

.. autofunction:: argmax