API Documentation
================

.. toctree::
   :maxdepth: 2

   attention
   functional
   models

Attention Modules
---------------

.. automodule:: jaxgarden.attention
   :members:
   :undoc-members:
   :show-inheritance:

Functional Interfaces
------------------

.. automodule:: jaxgarden.functional
   :members:
   :undoc-members:
   :show-inheritance: 

Models
------------------

.. automodule:: jaxgarden.models
   :members:
   :undoc-members:
   :show-inheritance: