optax.tree\_utils.tree\_batch\_shape
====================================

.. currentmodule:: optax.tree_utils

.. autofunction:: tree_batch_shape