Python reference
================

Core transforms
---------------
.. automodule:: jaxsaw.core
   :members:

Module composition
------------------
.. automodule:: jaxsaw.nn.module
   :members:

Layers
------
.. automodule:: jaxsaw.nn.layers
   :members:

Recurrent layers
----------------
.. automodule:: jaxsaw.nn.recurrent
   :members:

Attention
---------
.. automodule:: jaxsaw.nn.attention
   :members:

Loss functions
--------------
.. automodule:: jaxsaw.nn.functional
   :members:

Optimizers
----------
.. automodule:: jaxsaw.optim
   :members:

Reference models
----------------
.. automodule:: jaxsaw.models
   :members:

Data and training
-----------------
.. automodule:: jaxsaw.data
   :members:
.. automodule:: jaxsaw.training
   :members:

Checkpoints and inspection
--------------------------
.. automodule:: jaxsaw.checkpoint
   :members:
.. automodule:: jaxsaw.inspection
   :members:
.. automodule:: jaxsaw.experiments
   :members: