flax\_extra package
===================

.. automodule:: flax_extra
   :members:
   :undoc-members:
   :show-inheritance:

Subpackages
-----------

.. toctree::
   :maxdepth: 1

   flax_extra.checkpoint
   flax_extra.combinator
   flax_extra.data
   flax_extra.evaluation
   flax_extra.layer
   flax_extra.model
   flax_extra.operator
   flax_extra.training

Submodules
----------

.. toctree::
   :maxdepth: 1

   flax_extra.batch
   flax_extra.console
   flax_extra.random
   flax_extra.util
   flax_extra.version