training
----------------------------

Experimental API. See the `NNX page <https://flax.readthedocs.io/en/latest/nnx/index.html>`__ for more details.

.. toctree::
  :maxdepth: 3

  metrics
  optimizer
  ema