Discretax - State Space Models in JAX

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

• easy to use
• fast
• modular

• 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

• [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. 🎉

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 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)!

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]

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:

uv run pre-commit install

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

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.

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