traceTorch
traceTorch is a PyTorch library for stateful recurrent layers, built primarily for spiking neural networks.
It gives you SNN, RNN, and SSM-style layers that behave like ordinary PyTorch modules: one tensor in, one tensor out, hidden states kept inside the layer. The difference is that those hidden states are still easy to manage. Inherit from tt.Model, call zero_states(), detach_states(), save_states(), load_states(), TTcompile(), or TTdecompile(), and traceTorch handles every traceTorch layer buried inside the model.
pip install tracetorch
import torch
from torch import nn
import tracetorch as tt
class Net(tt.Model):
def __init__(self):
super().__init__()
self.net = nn.Sequential(
nn.Flatten(),
nn.Linear(784, 128),
tt.snn.LIB(num_neurons=128),
nn.Linear(128, 10),
tt.snn.LI(num_neurons=10),
)
def forward(self, x):
return self.net(x)
model = Net()
model.zero_states()
out = model(torch.rand(32, 1, 28, 28))
Why traceTorch?
- Hidden states stay hidden. Layers own their states, so model code stays readable.
- State management is explicit. Reset between sequences with
zero_states(), truncate history withdetach_states(), and save/load hidden states when needed. - SNNs are first-class.
tt.snncontains 32 leaky-integrator-based layers with binary, ternary, scaled ternary, continuous, dual, synaptic, and recurrent variants. - PyTorch composition stays normal. Put traceTorch layers inside
nn.Sequential, CNNs, MLPs, and custom PyTorch modules. - Feature dimensions are configurable. Use
dim=-1for MLP features,dim=-3for image channels, or any other target dimension. - Parameters can be scalar, per-neuron, fixed, learnable, or tensor-initialized.
Layer Families
| Module | Layers |
|---|---|
tt.snn |
LI, LIB, LIT, LITS families, including dual (D), synaptic (S), recurrent (R), and combined variants |
tt.rnn |
SimpleRNN, LSTM, GRU |
tt.ssm |
S4, S5, S6, Mamba adapted to traceTorch's one-timestep recurrent interface |
traceTorch's main focus is SNN experimentation. The RNN and SSM layers exist because the same state-management design is useful there too, but the SSM implementations are not meant to replace the official optimized sequence-parallel versions.
Documentation
Read the full documentation at https://yegor-men.github.io/tracetorch/.
Recommended path:
- Installation: install the package or editable repository.
- Quickstart: build and train a minimal traceTorch model.
- Introduction: understand the state model, SNN naming scheme, and design choices.
- Examples: follow MNIST and Heidelberg Digits examples from
examples/. - Tutorials: learn saving/loading states, compiling/decompiling, and custom layer creation.
- Reference: inspect API docstrings.
Examples
Runnable examples live in examples/.
git clone https://github.com/Yegor-men/tracetorch.git
cd tracetorch
pip install -e .
cd examples/mnist
pip install -r requirements.txt
python rate_coded.py
Current examples:
examples/mnist/rate_coded.py: rate-coded MNIST over repeated Bernoulli timesteps.examples/mnist/sequential.py: MNIST as a patch sequence.examples/mnist/noisy.py: MNIST with noisy repeated observations.examples/heidelberg_digits/main.py: Spiking Heidelberg Digits with Tonic.
These examples are written to show traceTorch mechanics clearly. They are not tuned as benchmark or SOTA training recipes.
Development Status
traceTorch is approaching its v1.0.0 release. The current focus is documentation, tests, examples, and API polish.
Author
Created by Yegor Menovchshikov.
License
MIT.
1 maintainer
Metadata
1 maintainer
Release files for tracetorch 0.19.2
For a detailed explanation of source distributions (sdists) and built distributions (wheels), please see the package formats documentation.
Source distribution (sdist)
| File | Size | Uploaded | |
|---|---|---|---|
| tracetorch-0.19.2.tar.gz | 38.6 kB | Details |
Built distribution (wheel)
| File | Interpreter | ABI | Platform | Reset |
|---|---|---|---|---|
| tracetorch-0.19.2-py3-none-any.whl | Python 3 | none | any | Details |
Total release size: 87.6 kB
Release files / tracetorch-0.19.2.tar.gz
| Download URL | tracetorch-0.19.2.tar.gz |
|---|---|
| Size | 38.6 kB |
| Tags | Source |
|
SHA-256 checksum How to use checksums |
480633ffcebce44aee53d4f798b802205c6e727783807ce8ab355588c6a9b44a
|
|
BLAKE2b-256 checksum How to use checksums |
87634cd4d35b91273ae613819f559be6d74dfe54b16e61b1139fc16c75b86f77
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.13.13
|
Release files / tracetorch-0.19.2-py3-none-any.whl
| Download URL | tracetorch-0.19.2-py3-none-any.whl |
|---|---|
| Size | 49.0 kB |
| Tags | Python 3 |
|
SHA-256 checksum How to use checksums |
6cc249b68f1eccf40f9c1d0fbf1abc2c77c62a9f08ca18489fc09fdc9d60e0ed
|
|
BLAKE2b-256 checksum How to use checksums |
cd2af4b89dff9d905abae52a492e7d1494dab050ce26c149ca09ccc4da057ad3
|
| Upload date | |
|
Uploaded using Trusted Publishing? What is trusted publishing? |
No |
| Uploaded via |
twine/6.2.0 CPython/3.13.13
|
