jax.lax.index\_in\_dim
======================

.. currentmodule:: jax.lax

.. autofunction:: index_in_dim