Linax - State Space Models in JAX

linax is a collection of state space models implemented in JAX. It is

• easy to use
• lightning-fast
• highly modular
• easily accessible.

• News
• Just get me Going
• Join the Community
• Installation
• Supported Models
• Contributing
• Core Contributors
• Citation

• [2025-10]: We are happy to officially launch the first version of linax. 🎉

If you don't care about the details, we provide example notebooks that are ready to use.

To join our growing community of JAX and state space model enthusiasts, join our Discord server. Feel free to write us a message (either there or to our personal email, see the bottom of this page) if you have any questions, comments, or just want to say hi!

🤫 Psssst! Rumor has it we are also developing an end-to-end JAX training pipeline. Stay tuned for JAX Lightning. So join the discord server to be the first to hear about our newest project(s)!

linax is available as a PyPI package. To install it via uv, just run

uv add linax

or

uv add linax[cu12]

If pip is your package manager of choice, run

pip install linax

or

pip install linax[cu12]

If you want to install the full library, especially if you want to contribute to the project, clone the linax repository and cd into it

git clone https://github.com/camail-official/linax.git
cd linax

If you want to install dependencies for CPU, run

uv sync

for GPU run

uv sync --extra cu12

To include development tooling (pre-commit, Ruff), install:

uv sync --extra dev

After installing the development dependencies (activate your environment if needed), enable the git hooks:

pre-commit install

If you want to contribute to the project, please check out contributing

This repository has been created and is maintained by:

• Benedict Armstrong
• Philipp Nazari
• Francesco Maria Ruscio

This work has been carried out within the Computational Applied Mathematics & AI Lab, led by T. Konstantin Rusch.

If you find this repository useful, please consider citing it.

@software{linax2025,
 title = {Linax: A Lightweight Collection of State Space Models in JAX},
 author = {Armstrong, Benedict and Nazari, Philipp and Ruscio, Francesco Maria},
 url = {https://github.com/camail-official/linax},
 year = {2025}
}