optax.trace
===========

.. currentmodule:: optax

.. autofunction:: trace