lite-mamba
A minimal, pure-TensorFlow implementation of Mamba with a multi-dilated causal depthwise conv front-end. No custom C++ or Triton kernels needed; works seamlessly on CPU, GPU, or TPU with standard TensorFlow ops.
📚 READ THE FULL DOCUMENTATION HERE 📚
Contains architecture details, API references, streaming inference guides, and overviews of the multi-branch variants (TFPTCNMamba, TFSTCNMamba, TFDPWCMamba).
Quick Start
Install
pip install lite-mamba
Basic Usage
from lite_mamba import TFPTCNMamba
import tensorflow as tf
x = tf.random.normal((2, 128, 512)) # (batch, seq, d_model)
m = TFPTCNMamba(d_model=512, d_conv=3, conv_dilations=(1, 2, 4, 8))
y = m(x)
print(y.shape) # (2, 128, 512)
Changelog
See CHANGELOG.md. 1.2.0 fixes a save_weights() failure
(Unable to synchronously create dataset (name already exists)) that affected every
model containing a Mamba block on Keras 2 / TF ≤ 2.15.
License
Apache-2.0
1 maintainer
Metadata
1 maintainer
Release files for lite-mamba 1.2.0
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| lite_mamba-1.2.0.tar.gz | 12.9 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| lite_mamba-1.2.0-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 21.3 kB
Release files / lite_mamba-1.2.0.tar.gz
| Download URL | lite_mamba-1.2.0.tar.gz |
|---|---|
| Size | 12.9 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
3e3223c546b272fde3f1c12e6f540aced9ccef7d1b1e926b36edc867c32b9c23
|
|
BLAKE2b-256 checksum How to use checksums |
547d16dab07b8f1e7eba1c19b38fa68214bb54109ecd258b9e55ffeba128e4db
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.11.15
|
Release files / lite_mamba-1.2.0-py3-none-any.whl
| Download URL | lite_mamba-1.2.0-py3-none-any.whl |
|---|---|
| Size | 8.4 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
f34e05074a7fc6e9f3627e4d43b8ccc0a31b6c0bd0060d34fbbabd86e09d7942
|
|
BLAKE2b-256 checksum How to use checksums |
87460f8bba90e1ad7c116f36ea0adc826e81892318d023155b0690267ed035cd
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/7.0.0 CPython/3.11.15
|
