discretax 0.2.1

State Space Models in Jax

Discretax - State Space Models in JAX

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

• easy to use
• fast
• modular

Table of contents

• Discretax - State Space Models in JAX
  • Table of contents
  • News
  • Just get me Going
  • Join the Community
  • Installation
    • Full Library Installation
  • Supported Models
  • Contributing
  • Core Contributors
  • Citation

News

• [2026-03]: After a big refactor, we are renaming the project from linax to discretax.
• [2025-10]: We are happy to launch the first beta 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

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

uv add discretax

or

uv add discretax[cu12]

If pip is your package manager of choice, run

pip install discretax

or

pip install discretax[cu12]

Full Library Installation

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

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

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	DeltaNet	Parallelizing Linear Transformers with the Delta Rule	sustcsonglin/flash-linear-attention	discretax
2024	LinOSS	Oscillatory State Space Models	tk-rusch/linoss	discretax
2023	LRU	Resurrecting Recurrent Neural Networks for Long Sequences	LRU paper	discretax
2022	S5	Simplified State Space Layers for Sequence Modeling	lindermanlab/S5	discretax
2022	S4D	On the Parameterization and Initialization of Diagonal State Space Models	state-spaces/s4	discretax

Contributing

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

Core Contributors

This repository has been created and is maintained by:

• Philipp Nazari
• Francesco Maria Ruscio
• Benedict Armstrong

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{discretax2025,
  title  = {Discretax: A Lightweight Collection of State Space Models in JAX},
  author = {Nazari, Philipp* and Ruscio, Francesco Maria* and Armstrong, Benedict and Rusch, T. Konstantin},
  url    = {https://github.com/camail-official/discretax},
  year   = {2025}
}