helpers
------------------------

.. automodule:: flax.nnx
.. currentmodule:: flax.nnx


.. autoclass:: Sequential
   :members:
.. autoclass:: List
   :members:
.. autoclass:: Dict
   :members:
.. autoclass:: TrainState
   :members: