jax.lax.axis\_index
===================

.. currentmodule:: jax.lax

.. autofunction:: axis_index