linax 0.1.0

State Space Models in Jax

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.

Table of contents

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

News

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

Just get me Going

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

Join the Community

To join our growing community of JAX and state space model enthusiasts, join our 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)!

Installation

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]

Full Library Installation

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

Supported Models

Year	Model	Paper	Code	Our implementation
2024	LinOSS	Oscillatory State Space Models	tk-rusch/linoss	linax
2023	LRU	Resurrecting Recurrent Neural Networks for Long Sequences	LRU paper	linax
2022	S5	Simplified State Space Layers for Sequence Modeling	lindermanlab/S5	linax
2022	S4D	On the Parameterization and Initialization of Diagonal State Space Models	state-spaces/s4	linax

Contributing

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

Core Contributors

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.

Citation

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}
}