lrnnx 1.0.0__tar.gz

lrnnx-1.0.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Society for Artificial Intelligence and Deep Learning
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

lrnnx-1.0.0/MANIFEST.in

@@ -0,0 +1 @@
+recursive-include csrc *.cpp *.cu *.cuh *.h

lrnnx-1.0.0/PKG-INFO

@@ -0,0 +1,183 @@
+Metadata-Version: 2.4
+Name: lrnnx
+Version: 1.0.0
+Summary: A library for Linear RNNs
+Author: SAiDl Team
+License-Expression: MIT
+Project-URL: Homepage, https://github.com/SforAiDl/lrnnx
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: numpy
+Requires-Dist: torch
+Requires-Dist: einops
+Requires-Dist: ninja
+Requires-Dist: packaging
+Requires-Dist: opt-einsum
+Requires-Dist: pykeops
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Requires-Dist: black; extra == "dev"
+Requires-Dist: isort; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Provides-Extra: conv1d
+Requires-Dist: causal-conv1d; extra == "conv1d"
+Dynamic: license-file
+
+<!---
+Copyright 2025 SAiDL Team. All rights reserved.
+
+Licensed under the MIT License; you may not use this file except in compliance
+with the License. You may obtain a copy of the License in the LICENSE file.
+-->
+
+# lrnnx: A library for Linear RNNs
+<p>
+	<a href="LICENSE"><img alt="License" src="https://img.shields.io/github/license/SforAiDl/lrnnx.svg?color=blue"></a>
+	<a href="https://arxiv.org/abs/2602.08810"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2602.08810-b31b1b.svg"></a>
+</p>
+
+A unified PyTorch library providing easy access to state-of-the-art Linear RNN architectures for sequence modeling.
+The technical report of this system was accepted to [EACL Student Research Workshop 2026](https://2026.eacl.org/calls/srw/).
+We recommend reading the report before using / contributing to the library.
+
+## Installation
+
+### From PyPI
+Since this library compiles custom CUDA kernels upon installation, we recommend using `--no-build-isolation` to avoid downloading a duplicate version of PyTorch during the build process.
+```bash
+# standard installation
+pip install lrnnx --no-build-isolation
+# with optional causal-conv1d
+pip install "lrnnx[conv1d]" --no-build-isolation
+# for development
+pip install "lrnnx[dev]" --no-build-isolation
+```
+
+### From Source
+We recommend installation with [`uv`](https://docs.astral.sh/uv/getting-started/installation/) for fast, reliable dependency management, though standard `pip` is fully supported.
+
+#### Using uv
+```bash
+git clone https://github.com/SforAiDl/lrnnx.git
+cd lrnnx
+# standard installation
+uv sync
+# with optional causal-conv1d
+uv sync --extra conv1d
+# for development
+uv sync --extra dev
+```
+
+#### Using pip
+```bash
+git clone https://github.com/SforAiDl/lrnnx.git
+cd lrnnx
+# standard installation
+pip install -e . --no-build-isolation
+# with optional causal-conv1d
+pip install -e ".[conv1d]" --no-build-isolation
+# for development
+pip install -e ".[dev]" --no-build-isolation
+```
+
+Note that since our library builds several custom CUDA kernels, it can take time for this installation to finish.
+Along with `causal-conv1d` the full installation can take about 30 minutes, depending on the number of CPUs available.
+
+## Model Zoo
+Our library provides implementations of the following Linear RNN architectures:
+- [S4](https://openreview.net/forum?id=uYLFoz1vlAC)
+- [S4D](https://dl.acm.org/doi/10.5555/3600270.3602877)
+- [S5](https://openreview.net/forum?id=Ai8Hw3AXqks)
+- [Event-SSM](https://www.computer.org/csdl/proceedings-article/icons/2024/686500a124/22lEawhJ0Va) (inside `S5`, use by passing `integration_timesteps`)
+- [LRU](https://dl.acm.org/doi/10.5555/3618408.3619518)
+- [S6](https://openreview.net/forum?id=tEYskw1VY2) (we implemented other discretizations)
+- [STREAM](https://arxiv.org/abs/2411.12603) (inside `S6`, use by passing `integration_timesteps`)
+- [RG-LRU](https://arxiv.org/abs/2402.19427)
+- [S7](https://arxiv.org/abs/2410.03464)
+- [aTENNuate](https://www.isca-archive.org/interspeech_2025/pei25_interspeech.html)
+
+We expose several levels of API for each model, including a scan, a recurrent step, and a full layer API matching the paper.
+For S5 we implement both a convolution based approach and a parallel scan approach.
+The latter is more stable and faster for most use cases, but the convolution based approach can be faster for very long sequences.
+
+## Usage
+
+### Training
+It is easy to instantiate a model from our library
+```python
+from lrnnx.models.lti import LRU
+from lrnnx.models.ltv import Mamba
+
+model_lti = LRU(d_model, d_state).cuda()
+x = torch.randn(
+	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
+)
+output = model_lti(x)
+
+model_ltv = Mamba(d_model, d_state).cuda()
+x = torch.randn(
+	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
+)
+output = model_ltv(x)
+```
+
+### Inference
+Linear RNNs in torch require special handling during inference, following [mamba](https://github.com/state-spaces/mamba), we also implement CUDA graphs based inference which reduces CPU overheads, this leads to > 10x speedup compared to using a simple for loop over the sequence length.
+The main file is [generation.py](lrnnx/generation.py) which provides a simple API for autoregressive generation with any of the models in our library.
+You can see a simple way to use it in our [benchmarking script](benchmarks/benchmark_inference.py).
+
+### Reproducing the Benchmarks from the paper
+This script will run both training and inference benchmarks.
+```bash
+python -m benchmarks.run_all
+```
+
+### Architectures
+We also implement some common architectures based on the models in our library, such as a U-Net (inspired from [aTENNuate](https://www.isca-archive.org/interspeech_2025/pei25_interspeech.html) ) and a hierarchical classifier (inspired from [Event-SSM](https://www.computer.org/csdl/proceedings-article/icons/2024/686500a124/22lEawhJ0Va)).
+Additionally, there is a [Language Model](lrnnx/models/language_model.py) architecture inspired from [Mamba](https://github.com/state-spaces/mamba) and [RG-LRU](https://arxiv.org/abs/2402.19427) which can be used for language modeling tasks, with replaceable LRNN and attention layers.
+This can be used as
+```python
+from lrnnx.models.language_model import LRNNLMHeadModel
+
+model = LRNNLMHeadModel(
+	d_model, d_state, num_layers, vocab_size, mixer_types=["s5", "s6", "attn"]
+)
+input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+logits = model(input_ids)
+```
+
+### Tutorial Overview
+
+Based on the architectures, there are tutorials on how to use them for 2 very popular use cases:
+1. [U-Net Seq2Seq for audio denoising Tutorial](tutorials/notebooks/01_UNet.ipynb)
+2. [Hierarchical Classification Tutorial](tutorials/notebooks/02_hierarchical_classifier.ipynb)
+
+## Contributing
+
+Please check out our [Contributing Guide](CONTRIBUTING.rst) for details on how to contribute to this project.
+
+## Citation
+
+If you use lrnnx in your research, please cite:
+
+```bibtex
+@misc{bania2026textttlrnnxlibrarylinearrnns,
+	title={$\texttt{lrnnx}$: A library for Linear RNNs}, 
+	author={Karan Bania and Soham Kalburgi and Manit Tanwar and Dhruthi and Aditya Nagarsekar and Harshvardhan Mestha and Naman Chibber and Raj Deshmukh and Anish Sathyanarayanan and Aarush Rathore and Pratham Chheda},
+	year={2026},
+	eprint={2602.08810},
+	archivePrefix={arXiv},
+	primaryClass={cs.LG},
+	url={https://arxiv.org/abs/2602.08810}, 
+}
+```
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Acknowledgments
+
+This library builds upon the excellent work of researchers who developed the individual LRNN models.
+Please see individual model documentation for proper citations of the original papers.

lrnnx-1.0.0/README.md

@@ -0,0 +1,157 @@
+<!---
+Copyright 2025 SAiDL Team. All rights reserved.
+
+Licensed under the MIT License; you may not use this file except in compliance
+with the License. You may obtain a copy of the License in the LICENSE file.
+-->
+
+# lrnnx: A library for Linear RNNs
+<p>
+	<a href="LICENSE"><img alt="License" src="https://img.shields.io/github/license/SforAiDl/lrnnx.svg?color=blue"></a>
+	<a href="https://arxiv.org/abs/2602.08810"><img alt="arXiv" src="https://img.shields.io/badge/arXiv-2602.08810-b31b1b.svg"></a>
+</p>
+
+A unified PyTorch library providing easy access to state-of-the-art Linear RNN architectures for sequence modeling.
+The technical report of this system was accepted to [EACL Student Research Workshop 2026](https://2026.eacl.org/calls/srw/).
+We recommend reading the report before using / contributing to the library.
+
+## Installation
+
+### From PyPI
+Since this library compiles custom CUDA kernels upon installation, we recommend using `--no-build-isolation` to avoid downloading a duplicate version of PyTorch during the build process.
+```bash
+# standard installation
+pip install lrnnx --no-build-isolation
+# with optional causal-conv1d
+pip install "lrnnx[conv1d]" --no-build-isolation
+# for development
+pip install "lrnnx[dev]" --no-build-isolation
+```
+
+### From Source
+We recommend installation with [`uv`](https://docs.astral.sh/uv/getting-started/installation/) for fast, reliable dependency management, though standard `pip` is fully supported.
+
+#### Using uv
+```bash
+git clone https://github.com/SforAiDl/lrnnx.git
+cd lrnnx
+# standard installation
+uv sync
+# with optional causal-conv1d
+uv sync --extra conv1d
+# for development
+uv sync --extra dev
+```
+
+#### Using pip
+```bash
+git clone https://github.com/SforAiDl/lrnnx.git
+cd lrnnx
+# standard installation
+pip install -e . --no-build-isolation
+# with optional causal-conv1d
+pip install -e ".[conv1d]" --no-build-isolation
+# for development
+pip install -e ".[dev]" --no-build-isolation
+```
+
+Note that since our library builds several custom CUDA kernels, it can take time for this installation to finish.
+Along with `causal-conv1d` the full installation can take about 30 minutes, depending on the number of CPUs available.
+
+## Model Zoo
+Our library provides implementations of the following Linear RNN architectures:
+- [S4](https://openreview.net/forum?id=uYLFoz1vlAC)
+- [S4D](https://dl.acm.org/doi/10.5555/3600270.3602877)
+- [S5](https://openreview.net/forum?id=Ai8Hw3AXqks)
+- [Event-SSM](https://www.computer.org/csdl/proceedings-article/icons/2024/686500a124/22lEawhJ0Va) (inside `S5`, use by passing `integration_timesteps`)
+- [LRU](https://dl.acm.org/doi/10.5555/3618408.3619518)
+- [S6](https://openreview.net/forum?id=tEYskw1VY2) (we implemented other discretizations)
+- [STREAM](https://arxiv.org/abs/2411.12603) (inside `S6`, use by passing `integration_timesteps`)
+- [RG-LRU](https://arxiv.org/abs/2402.19427)
+- [S7](https://arxiv.org/abs/2410.03464)
+- [aTENNuate](https://www.isca-archive.org/interspeech_2025/pei25_interspeech.html)
+
+We expose several levels of API for each model, including a scan, a recurrent step, and a full layer API matching the paper.
+For S5 we implement both a convolution based approach and a parallel scan approach.
+The latter is more stable and faster for most use cases, but the convolution based approach can be faster for very long sequences.
+
+## Usage
+
+### Training
+It is easy to instantiate a model from our library
+```python
+from lrnnx.models.lti import LRU
+from lrnnx.models.ltv import Mamba
+
+model_lti = LRU(d_model, d_state).cuda()
+x = torch.randn(
+	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
+)
+output = model_lti(x)
+
+model_ltv = Mamba(d_model, d_state).cuda()
+x = torch.randn(
+	batch_size, seq_len, d_model, dtype=torch.float32, device="cuda"
+)
+output = model_ltv(x)
+```
+
+### Inference
+Linear RNNs in torch require special handling during inference, following [mamba](https://github.com/state-spaces/mamba), we also implement CUDA graphs based inference which reduces CPU overheads, this leads to > 10x speedup compared to using a simple for loop over the sequence length.
+The main file is [generation.py](lrnnx/generation.py) which provides a simple API for autoregressive generation with any of the models in our library.
+You can see a simple way to use it in our [benchmarking script](benchmarks/benchmark_inference.py).
+
+### Reproducing the Benchmarks from the paper
+This script will run both training and inference benchmarks.
+```bash
+python -m benchmarks.run_all
+```
+
+### Architectures
+We also implement some common architectures based on the models in our library, such as a U-Net (inspired from [aTENNuate](https://www.isca-archive.org/interspeech_2025/pei25_interspeech.html) ) and a hierarchical classifier (inspired from [Event-SSM](https://www.computer.org/csdl/proceedings-article/icons/2024/686500a124/22lEawhJ0Va)).
+Additionally, there is a [Language Model](lrnnx/models/language_model.py) architecture inspired from [Mamba](https://github.com/state-spaces/mamba) and [RG-LRU](https://arxiv.org/abs/2402.19427) which can be used for language modeling tasks, with replaceable LRNN and attention layers.
+This can be used as
+```python
+from lrnnx.models.language_model import LRNNLMHeadModel
+
+model = LRNNLMHeadModel(
+	d_model, d_state, num_layers, vocab_size, mixer_types=["s5", "s6", "attn"]
+)
+input_ids = torch.randint(0, vocab_size, (batch_size, seq_len))
+logits = model(input_ids)
+```
+
+### Tutorial Overview
+
+Based on the architectures, there are tutorials on how to use them for 2 very popular use cases:
+1. [U-Net Seq2Seq for audio denoising Tutorial](tutorials/notebooks/01_UNet.ipynb)
+2. [Hierarchical Classification Tutorial](tutorials/notebooks/02_hierarchical_classifier.ipynb)
+
+## Contributing
+
+Please check out our [Contributing Guide](CONTRIBUTING.rst) for details on how to contribute to this project.
+
+## Citation
+
+If you use lrnnx in your research, please cite:
+
+```bibtex
+@misc{bania2026textttlrnnxlibrarylinearrnns,
+	title={$\texttt{lrnnx}$: A library for Linear RNNs}, 
+	author={Karan Bania and Soham Kalburgi and Manit Tanwar and Dhruthi and Aditya Nagarsekar and Harshvardhan Mestha and Naman Chibber and Raj Deshmukh and Anish Sathyanarayanan and Aarush Rathore and Pratham Chheda},
+	year={2026},
+	eprint={2602.08810},
+	archivePrefix={arXiv},
+	primaryClass={cs.LG},
+	url={https://arxiv.org/abs/2602.08810}, 
+}
+```
+
+## License
+
+This project is licensed under the MIT License - see the [LICENSE](LICENSE) file for details.
+
+## Acknowledgments
+
+This library builds upon the excellent work of researchers who developed the individual LRNN models.
+Please see individual model documentation for proper citations of the original papers.

lrnnx-1.0.0/benchmarks/__init__.py

@@ -0,0 +1 @@
+"""Benchmark utilities for lrnnx models."""

lrnnx-1.0.0/benchmarks/benchmark_cauchy.py

@@ -0,0 +1,62 @@
+import math
+from functools import partial
+
+import torch
+from benchmark.utils import (
+    benchmark_all,
+    benchmark_backward,
+    benchmark_combined,
+    benchmark_forward,
+)
+from einops import rearrange
+
+from .cauchy import cauchy_mult, cauchy_mult_keops, cauchy_mult_torch
+
+
+def generate_data(batch_size, N, L, symmetric=True, device="cuda"):
+    if not symmetric:
+        v = torch.randn(
+            batch_size,
+            N,
+            dtype=torch.complex64,
+            device=device,
+            requires_grad=True,
+        )
+        w = torch.randn(
+            batch_size,
+            N,
+            dtype=torch.complex64,
+            device=device,
+            requires_grad=True,
+        )
+        z = torch.randn(L, dtype=torch.complex64, device=device)
+    else:
+        assert N % 2 == 0
+        v_half = torch.randn(
+            batch_size, N // 2, dtype=torch.complex64, device=device
+        )
+        v = torch.cat([v_half, v_half.conj()], dim=-1).requires_grad_(True)
+        w_half = torch.randn(
+            batch_size, N // 2, dtype=torch.complex64, device=device
+        )
+        w = torch.cat([w_half, w_half.conj()], dim=-1).requires_grad_(True)
+        z = torch.exp(1j * torch.randn(L, dtype=torch.float32, device=device))
+    return v, z, w
+
+
+if __name__ == "__main__":
+    device = "cuda"
+    bs = 1024
+    N = 64
+    L = 16384
+
+    v, z, w = generate_data(bs, N, L, symmetric=True)
+    v_half = v[:, : N // 2].clone().detach().requires_grad_(True)
+    w_half = w[:, : N // 2].clone().detach().requires_grad_(True)
+
+    repeat = 30
+    benchmark_all(repeat, cauchy_mult_keops, v, z, w, desc="Cauchy mult keops")
+    fn = partial(cauchy_mult, symmetric=False)
+    benchmark_all(repeat, fn, v, z, w, desc="Cauchy mult")
+    fn = partial(cauchy_mult, symmetric=True)
+    benchmark_all(repeat, fn, v_half, z, w_half, desc="Cauchy mult symmetric")

lrnnx-1.0.0/benchmarks/benchmark_cauchy_tune.py

@@ -0,0 +1,88 @@
+import argparse
+import importlib
+import json
+
+import torch
+from benchmark.utils import benchmark_forward
+
+
+def generate_data(batch_size, N, L, symmetric=True, device="cuda"):
+    if not symmetric:
+        v = torch.randn(
+            batch_size,
+            N,
+            dtype=torch.complex64,
+            device=device,
+            requires_grad=True,
+        )
+        w = torch.randn(
+            batch_size,
+            N,
+            dtype=torch.complex64,
+            device=device,
+            requires_grad=True,
+        )
+        z = torch.randn(L, dtype=torch.complex64, device=device)
+    else:
+        assert N % 2 == 0
+        v_half = torch.randn(
+            batch_size, N // 2, dtype=torch.complex64, device=device
+        )
+        v = torch.cat([v_half, v_half.conj()], dim=-1).requires_grad_(True)
+        w_half = torch.randn(
+            batch_size, N // 2, dtype=torch.complex64, device=device
+        )
+        w = torch.cat([w_half, w_half.conj()], dim=-1).requires_grad_(True)
+        z = torch.exp(1j * torch.randn(L, dtype=torch.float32, device=device))
+    return v, z, w
+
+
+parser = argparse.ArgumentParser(description="Tuning Cauchy multiply")
+parser.add_argument("--name", default="cauchy_mult")
+parser.add_argument(
+    "--mode", default="forward", choices=["forward", "backward"]
+)
+parser.add_argument("-bs", "--batch-size", default=1024, type=int)
+parser.add_argument("-N", default=64, type=int)
+parser.add_argument("-L", default=2**14, type=int)
+
+
+if __name__ == "__main__":
+    args = parser.parse_args()
+    device = "cuda"
+    bs = args.batch_size
+    N = args.N
+    L = args.L
+    repeat = 30
+    v, z, w = generate_data(bs, N, L, symmetric=True)
+    v_half = v[:, : N // 2].clone().detach().requires_grad_(True)
+    w_half = w[:, : N // 2].clone().detach().requires_grad_(True)
+
+    tuning_extension_name = args.name
+    # print('Extension name:', tuning_extension_name)
+    module = importlib.import_module(tuning_extension_name)
+    if args.mode == "forward":
+        _, m = benchmark_forward(
+            repeat,
+            module.cauchy_mult_sym_fwd,
+            v_half,
+            z,
+            w_half,
+            verbose=False,
+            desc="Cauchy mult symmetric fwd",
+        )
+    else:
+        out = module.cauchy_mult_sym_fwd(v_half, z, w_half)
+        dout = torch.randn_like(out)
+        _, m = benchmark_forward(
+            repeat,
+            module.cauchy_mult_sym_bwd,
+            v_half,
+            z,
+            w_half,
+            dout,
+            verbose=False,
+            desc="Cauchy mult symmetric bwd",
+        )
+    result_dict = dict(time_mean=m.mean, time_iqr=m.iqr)
+    print(json.dumps(result_dict))