Skip to content

Commit 3440a0d

Browse files
phnazariphnazari
authored
Restructure (#66)
* major refactoring. No more configs * Refactor: Rename classes to Abstract* for clarity * moved example and better model printing * removed all references to linax and replaced them by discretax * removed examples and updated names * update docs --------- Co-authored-by: Benedict Armstrong, Philipp Nazari, Francesco Ruscio
1 parent e465db4 commit 3440a0d

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

95 files changed

+2535
-2984
lines changed

‎.github/CONTRIBUTING.md‎

Lines changed: 9 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
1-
# Contributing to Linax
1+
# Contributing to Discretax
22

3-
Thank you for your interest in contributing to Linax! 🎉
3+
Thank you for your interest in contributing to Discretax!
44

55
We welcome contributions of all kinds: bug reports, feature requests, documentation improvements, and code contributions.
66

@@ -34,8 +34,8 @@ Before contributing, please:
3434
1. **Fork and clone the repository:**
3535

3636
```bash
37-
git clone https://github.com/YOUR_USERNAME/linax.git
38-
cd linax
37+
git clone https://github.com/YOUR_USERNAME/discretax.git
38+
cd discretax
3939
```
4040

4141
2. **Install dependencies (including dev tools):**
@@ -123,7 +123,7 @@ uv run pytest
123123
uv run pytest tests/test_smoke.py
124124

125125
# Run with coverage
126-
uv run pytest --cov=linax
126+
uv run pytest --cov=discretax
127127
```
128128

129129
### Writing Tests
@@ -193,19 +193,19 @@ changes
193193
## Community
194194

195195
- **Discord**: [Join our Discord server](https://discord.gg/VazrGCxeT7)
196-
- **Issues**: [GitHub Issues](https://github.com/camail-official/linax/issues)
197-
- **Website**: [camail.org/linax](https://camail.org/linax/)
196+
- **Issues**: [GitHub Issues](https://github.com/camail-official/discretax/issues)
197+
- **Website**: [camail.org/discretax](https://camail.org/discretax/)
198198

199199
---
200200

201201
## Questions?
202202

203203
If you have questions about contributing, feel free to:
204204

205-
- Open a [GitHub Discussion](https://github.com/camail-official/linax/discussions)
205+
- Open a [GitHub Discussion](https://github.com/camail-official/discretax/discussions)
206206
- Ask in our [Discord server](https://discord.gg/VazrGCxeT7)
207207
- Email the maintainers (see `pyproject.toml` for contact info)
208208

209209
---
210210

211-
Thank you for contributing to Linax! 🚀
211+
Thank you for contributing to Discretax!

‎.github/ISSUE_TEMPLATE/bug_report.yml‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
name: 🐛 Bug Report
2-
description: File a bug report to help us improve Linax
2+
description: File a bug report to help us improve Discretax
33
title: "[Bug]: "
44
labels: ["bug", "triage"]
55
body:
66
- type: markdown
77
attributes:
88
value: |
9-
Thanks for taking the time to fill out this bug report! 🙏
9+
Thanks for taking the time to fill out this bug report!
1010
1111
- type: textarea
1212
id: description
@@ -23,7 +23,7 @@ body:
2323
label: Steps to Reproduce
2424
description: How can we reproduce this issue?
2525
placeholder: |
26-
1. Import linax
26+
1. Import discretax
2727
2. Create a model with '...'
2828
3. Run '...'
2929
4. See error
@@ -38,7 +38,7 @@ body:
3838
render: python
3939
placeholder: |
4040
import jax
41-
import linax
41+
import discretax
4242
4343
# Your code here...
4444
validations:
@@ -58,11 +58,11 @@ body:
5858
id: environment
5959
attributes:
6060
label: Environment
61-
description: Python version, JAX version, Linax version, OS, hardware (GPU/CPU/TPU)
61+
description: Python version, JAX version, Discretax version, OS, hardware (GPU/CPU/TPU)
6262
placeholder: |
6363
Python 3.12.0
6464
JAX 0.4.23
65-
Linax main branch
65+
Discretax main branch
6666
macOS 14.0
6767
CPU
6868
validations:

‎.github/ISSUE_TEMPLATE/config.yml‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,10 @@ blank_issues_enabled: true
22
contact_links:
33
- name: 💬 Discord Community
44
url: https://discord.gg/VazrGCxeT7
5-
about: Ask questions and discuss Linax with the community
5+
about: Ask questions and discuss Discretax with the community
66
- name: 📚 Documentation
7-
url: https://camail.org/linax/
7+
url: https://camail.org/discretax/
88
about: Check out the documentation for guides and examples
99
- name: ❓ GitHub Discussions
10-
url: https://github.com/camail-official/linax/discussions
11-
about: Ask questions or start a discussion about Linax
10+
url: https://github.com/camail-official/discretax/discussions
11+
about: Ask questions or start a discussion about Discretax

‎.github/ISSUE_TEMPLATE/feature_request.yml‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,12 +1,12 @@
11
name: ✨ Feature Request
2-
description: Suggest a new feature or enhancement for Linax
2+
description: Suggest a new feature or enhancement for Discretax
33
title: "[Feature]: "
44
labels: ["enhancement"]
55
body:
66
- type: markdown
77
attributes:
88
value: |
9-
Thanks for suggesting a new feature! 💡
9+
Thanks for suggesting a new feature!
1010
1111
- type: textarea
1212
id: problem
@@ -22,7 +22,7 @@ body:
2222
attributes:
2323
label: Proposed Solution
2424
description: Describe how you'd like this to work
25-
placeholder: I would like Linax to...
25+
placeholder: I would like Discretax to...
2626
validations:
2727
required: true
2828

@@ -33,10 +33,10 @@ body:
3333
description: Show us how you'd like to use this feature
3434
render: python
3535
placeholder: |
36-
import linax
36+
import discretax
3737
3838
# Example of how the feature would be used
39-
model = linax.MyNewFeature(...)
39+
model = discretax.MyNewFeature(...)
4040
validations:
4141
required: false
4242

‎.gitignore‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -79,4 +79,4 @@ target/
7979

8080
# Miscellaneous
8181
*.DS_Store
82-
*/data_dir/
82+
data_dir/

‎Makefile‎

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -55,3 +55,6 @@ clean-env: ## Remove virtual environment and lockfile
5555
remake-lockfile: clean-env ## Recreate lockfile and environment from scratch
5656
uv sync --extra dev --extra examples --extra docs
5757
@echo "✅ Fresh environment created with new lockfile"
58+
59+
build-docs: ## Build docs
60+
uv run mkdocs build

‎README.md‎

Lines changed: 62 additions & 41 deletions
Original file line numberDiff line numberDiff line change
@@ -1,43 +1,47 @@
1-
# Linax - State Space Models in JAX
2-
31
<div align="center">
4-
<img alt="Linax Banner" src="https://raw.githubusercontent.com/camail-official/linax/refs/heads/main/assets/logo.png" style="padding-bottom: 2rem;" />
5-
</div>
62

7-
<div align="center">
3+
# Discretax - State Space Models in JAX
4+
5+
<img src="https://raw.githubusercontent.com/camail-official/discretax/main/assets/logo.png" alt="Discretax logo" width="200"/>
6+
87

98
[![pre-commit](https://img.shields.io/badge/Pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://github.com/pre-commit/pre-commit)
10-
[![tests](https://github.com/camail-official/linax/actions/workflows/tests.yml/badge.svg)](https://github.com/camail-official/linax/actions/workflows/test.yaml)
11-
[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/camail-official/linax/blob/main/LICENSE)
9+
[![tests](https://github.com/camail-official/discretax/actions/workflows/tests.yml/badge.svg)](https://github.com/camail-official/discretax/actions/workflows/test.yaml)
10+
[![license](https://img.shields.io/badge/License-MIT-green.svg?labelColor=gray)](https://github.com/camail-official/discretax/blob/main/LICENSE)
1211
[![uv](https://img.shields.io/endpoint?url=https://raw.githubusercontent.com/astral-sh/uv/main/assets/badge/v0.json)](https://github.com/astral-sh/uv)[![JAX](https://img.shields.io/badge/JAX-0.7%2B-0A7AAA?logo=jax&logoColor=white)](https://github.com/google/jax)
13-
[![PyPI version](https://img.shields.io/pypi/v/linax)](https://pypi.org/project/linax/)
12+
[![PyPI version](https://img.shields.io/pypi/v/discretax)](https://pypi.org/project/discretax/)
1413
[![Discord](https://img.shields.io/badge/Discord-%235865F2.svg?logo=discord&logoColor=white)](https://discord.gg/VazrGCxeT7)
1514

1615
</div>
1716

18-
[linax](https://github.com/camail-official/linax) is a collection of state space models implemented in JAX. It is
17+
[discretax](https://github.com/camail-official/discretax) is a collection of state space models implemented in JAX. It is
1918

2019
- easy to use
21-
- lightning-fast
22-
- highly modular
23-
- easily accessible.
20+
- fast
21+
- modular
2422

2523
## Table of contents
26-
- [News](#news)
27-
- [Just get me Going](#just-get-me-going)
28-
- [Join the Community](#join-the-community)
29-
- [Installation](#installation)
30-
- [Supported Models](#supported-models)
31-
- [Contributing](#contributing)
32-
- [Core Contributors](#core-contributors)
33-
- [Citation](#citation)
24+
25+
- [Discretax - State Space Models in JAX](#discretax---state-space-models-in-jax)
26+
- [Table of contents](#table-of-contents)
27+
- [News](#news)
28+
- [Just get me Going](#just-get-me-going)
29+
- [Join the Community](#join-the-community)
30+
- [Installation](#installation)
31+
- [Full Library Installation](#full-library-installation)
32+
- [Supported Models](#supported-models)
33+
- [Contributing](#contributing)
34+
- [Core Contributors](#core-contributors)
35+
- [Citation](#citation)
3436

3537
## News
36-
- [2025-10]: We are happy to officially launch the first version of linax. 🎉
38+
39+
- [2026-03]: After a big refactor, we are renaming the project from linax to discretax.
40+
- [2025-10]: We are happy to launch the first beta version of linax. 🎉
3741

3842
## Just get me Going
39-
If you don't care about the details, we provide [example notebooks](docs/examples/) that are ready to use.
4043

44+
If you don't care about the details, we provide [example notebooks](examples/) that are ready to use.
4145

4246
## Join the Community
4347

@@ -46,80 +50,97 @@ To join our growing community of JAX and state space model enthusiasts, join our
4650
🤫 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)!
4751

4852
## Installation
49-
[linax](https://github.com/camail-official/linax) is available as a PyPI package. To install it via uv, just run
53+
54+
[discretax](https://github.com/camail-official/discretax) is available as a PyPI package. To install it via uv, just run
55+
5056
```bash
51-
uv add linax
57+
uv add discretax
5258
```
59+
5360
or
61+
5462
```bash
55-
uv add linax[cu12]
63+
uv add discretax[cu12]
5664
```
5765

5866
If pip is your package manager of choice, run
67+
5968
```bash
60-
pip install linax
69+
pip install discretax
6170
```
71+
6272
or
73+
6374
```bash
64-
pip install linax[cu12]
75+
pip install discretax[cu12]
6576
```
6677

6778
### Full Library Installation
68-
If you want to install the full library, especially if you want to **contribute** to the project, clone the [linax](https://github.com/camail-official/linax) repository and cd into it
79+
80+
If you want to install the full library, especially if you want to **contribute** to the project, clone the [discretax](https://github.com/camail-official/discretax) repository and cd into it
81+
6982
```bash
70-
git clone https://github.com/camail-official/linax.git
71-
cd linax
83+
git clone https://github.com/camail-official/discretax.git
84+
cd discretax
7285
```
7386

7487
If you want to install dependencies for CPU, run
88+
7589
```bash
7690
uv sync
7791
```
92+
7893
for GPU run
94+
7995
```bash
8096
uv sync --extra cu12
8197
```
8298

8399
To include development tooling (pre-commit, Ruff), install:
100+
84101
```bash
85102
uv sync --extra dev
86103
```
104+
87105
After installing the development dependencies (activate your environment if needed), enable the git hooks:
106+
88107
```bash
89108
pre-commit install
90109
```
91110

92111
## Supported Models
93112

94-
| Year | Model | Paper | Code | Our implementation |
95-
|------|-------|-------|------|---------------------|
96-
| 2024 | LinOSS | [Oscillatory State Space Models](https://openreview.net/pdf?id=GRMfXcAAFh) | [tk-rusch/linoss](https://github.com/tk-rusch/linoss) | [linax](https://github.com/camail-official/linax/blob/main/src/linax/models/linoss.py) |
97-
| 2023 | LRU | [Resurrecting Recurrent Neural Networks for Long Sequences](https://proceedings.mlr.press/v202/orvieto23a/orvieto23a.pdf) | [LRU paper](https://proceedings.mlr.press/v202/orvieto23a/orvieto23a.pdf) | [linax](https://github.com/camail-official/linax/blob/main/src/linax/models/lru.py) |
98-
| 2022 | S5 | [Simplified State Space Layers for Sequence Modeling](https://openreview.net/pdf?id=Ai8Hw3AXqks) | [lindermanlab/S5](https://github.com/lindermanlab/S5) | [linax](https://github.com/camail-official/linax/blob/main/src/linax/models/s5.py) |
99-
| 2022 | S4D | [On the Parameterization and Initialization of Diagonal State Space Models](https://proceedings.neurips.cc/paper_files/paper/2022/file/e9a32fade47b906de908431991440f7c-Paper-Conference.pdf) | [state-spaces/s4](https://github.com/state-spaces/s4) | [linax](https://github.com/camail-official/linax/blob/main/src/linax/sequence_mixers/s4d.py) |
113+
| Year | Model | Paper | Code | Our implementation |
114+
| ---- | ------ | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------- | -------------------------------------------------------------------------------------------------------- |
115+
| 2024 | LinOSS | [Oscillatory State Space Models](https://openreview.net/pdf?id=GRMfXcAAFh) | [tk-rusch/linoss](https://github.com/tk-rusch/linoss) | [discretax](https://github.com/camail-official/discretax/blob/main/src/discretax/models/linoss.py) |
116+
| 2023 | LRU | [Resurrecting Recurrent Neural Networks for Long Sequences](https://proceedings.mlr.press/v202/orvieto23a/orvieto23a.pdf) | [LRU paper](https://proceedings.mlr.press/v202/orvieto23a/orvieto23a.pdf) | [discretax](https://github.com/camail-official/discretax/blob/main/src/discretax/models/lru.py) |
117+
| 2022 | S5 | [Simplified State Space Layers for Sequence Modeling](https://openreview.net/pdf?id=Ai8Hw3AXqks) | [lindermanlab/S5](https://github.com/lindermanlab/S5) | [discretax](https://github.com/camail-official/discretax/blob/main/src/discretax/models/s5.py) |
118+
| 2022 | S4D | [On the Parameterization and Initialization of Diagonal State Space Models](https://proceedings.neurips.cc/paper_files/paper/2022/file/e9a32fade47b906de908431991440f7c-Paper-Conference.pdf) | [state-spaces/s4](https://github.com/state-spaces/s4) | [discretax](https://github.com/camail-official/discretax/blob/main/src/discretax/sequence_mixers/s4d.py) |
100119

101120
## Contributing
121+
102122
If you want to contribute to the project, please check out [contributing](docs/contributing.md)
103123

104124
## Core Contributors
105125

106126
This repository has been created and is maintained by:
107127

108-
- [Benedict Armstrong](https://github.com/benedict-armstrong)
109128
- [Philipp Nazari](https://phnazari.github.io)
110129
- [Francesco Maria Ruscio](https://github.com/francescoshox)
130+
- [Benedict Armstrong](https://github.com/benedict-armstrong)
111131

112132
This work has been carried out within the [Computational Applied Mathematics & AI Lab](https://camail.org),
113133
led by [T. Konstantin Rusch](https://github.com/tk-rusch).
114134

115135
## Citation
136+
116137
If you find this repository useful, please consider citing it.
117138

118139
```bib
119-
@software{linax2025,
120-
title = {Linax: A Lightweight Collection of State Space Models in JAX},
121-
author = {Armstrong, Benedict and Nazari, Philipp and Ruscio, Francesco Maria},
122-
url = {https://github.com/camail-official/linax},
140+
@software{discretax2025,
141+
title = {Discretax: A Lightweight Collection of State Space Models in JAX},
142+
author = {Nazari, Philipp* and Ruscio, Francesco Maria* and Armstrong, Benedict and Rusch, T. Konstantin},
143+
url = {https://github.com/camail-official/discretax},
123144
year = {2025}
124145
}
125146
```

‎assets/logo.png‎

130 KB
Loading

0 commit comments

Comments
 (0)