.. highlight:: python

Reference
=========

.. automodule:: powerpax

Vectorization
-------------

.. autofunction:: powerpax.chunked_vmap

Loops
-----

.. autofunction:: powerpax.sliced_scan

.. autofunction:: powerpax.checkpoint_chunked_scan

Pytrees
-------

.. autoclass:: powerpax.Static