Skip to content

Getting Started¤

Discretax logo

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

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

Table of contents¤

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

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

Contributing¤

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

Core Contributors¤

This repository has been created and is maintained by:

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