optax.microbatching.reshape\_batch\_axis
========================================

.. currentmodule:: optax.microbatching

.. autofunction:: reshape_batch_axis