API Reference
=============

.. toctree::
   :maxdepth: 4

   flax.nnx/index
   flax.core.frozen_dict
   flax.struct
   flax.training