Flax fundamentals
=================

.. toctree::
   :maxdepth: 1

   JAX 101 <https://jax.readthedocs.io/en/latest/jax-101/index.html>
   flax_basics
   state_params
   setup_or_nncompact
   arguments