torchtsetlin

GPU-enabled, PyTorch-native Tsetlin machines — models, data preparation, training, evaluation, interpretation and visualisation in one package.

A Tsetlin machine learns propositional rules (IF x0 AND NOT x1 THEN class 1) with teams of learning automata instead of gradients. torchtsetlin implements the family as ordinary torch.nn.Modules: the automata live in tensors, so a model moves between CPU and GPU with .to(device), is saved with state_dict(), and learns from mini-batches with a single model.update(x, y) call.

import torchtsetlin as tt

x_train, y_train = tt.data.make_noisy_xor(5000, noise=0.4)          # Boolean features, int labels
x_test, y_test = tt.data.make_noisy_xor(2000, noise=0.0, seed=1)

model = tt.TsetlinMachine(n_features=12, n_classes=2, n_clauses=20, T=15, s=3.9).to("cuda")

# PyTorch-style loop: update() is the analogue of loss.backward(); optimizer.step()
for epoch in range(50):
    model.train()
    for i in range(0, 5000, 10):
        model.update(x_train[i:i+10].cuda(), y_train[i:i+10].cuda())
    model.eval()
    acc = (model(x_test.cuda()).argmax(1) == y_test.cuda()).float().mean()

print(model.rules()[:2])   # ['IF x0 AND NOT x1 THEN 1', 'IF x1 AND NOT x0 THEN 1']

Or with the built-in trainer:

trainer = tt.Trainer(model, batch_size=10, callbacks=[tt.EarlyStopping(patience=20)])
trainer.fit((x_train, y_train), epochs=100, val_data=(x_test, y_test))
trainer.evaluate((x_test, y_test))          # {'accuracy': 0.99}

Features

Installation

torchtsetlin is published on PyPI: https://pypi.org/project/torchtsetlin/. It needs Python 3.9+ and PyTorch 2.0+.

1. Install PyTorch first, choosing the CUDA or CPU build for your machine from pytorch.org, e.g.

   pip install torch                                                     # default build
   pip install torch --index-url https://download.pytorch.org/whl/cpu    # CPU-only
   

2. Install torchtsetlin from PyPI:

   pip install torchtsetlin
   

3. Optional extras:

   pip install "torchtsetlin[viz]"      # matplotlib for torchtsetlin.viz
   pip install "torchtsetlin[vision]"   # torchvision for the MNIST / torchvision / segmentation dataset helpers
   pip install "torchtsetlin[sklearn]"  # scikit-learn interop in the examples
   pip install "torchtsetlin[all]"      # everything, including docs and dev tools
   

Check the installation:

import torch, torchtsetlin as tt
print(tt.__version__, torch.cuda.is_available())

From source

git clone https://github.com/vlbthambawita/torchtsetlin
cd torchtsetlin
pip install -e ".[dev]"
pytest            # run the test-suite (uses the GPU when available)
mkdocs serve      # browse the documentation locally

How learning works (in one paragraph)

Each clause is an AND of literals (x_k or NOT x_k); every (clause, literal) pair has a Tsetlin automaton whose integer state decides whether the literal is included. For a training example, clauses of the true class receive Type I feedback (memorise the example's True literals, forget the rest) and clauses of a random other class receive Type II feedback (add a False literal so the clause stops matching), each with a probability controlled by the vote margin T. torchtsetlin evaluates a mini-batch with one matrix product, counts the feedback events per (clause, literal) with another, and turns the counts into binomial state transitions — the whole update is a handful of tensor ops. See the concept pages for details and references.

Documentation

• Getting started, concepts, guides, examples and the full API reference: https://vlbthambawita.github.io/torchtsetlin/ (or mkdocs serve locally).
• Measured CPU vs GPU performance (throughput, scaling, memory, where the time goes): Benchmarks -> CPU vs GPU; the harness lives in benchmarks/.
• Runnable scripts in examples/ (see its README): Noisy XOR, tabular data with thermometer encoding, regression, MNIST flat and convolutional, four convolutional walkthroughs covering 2-D clauses (shapes_conv.py), 1-D signals (conv1d_ramps.py), convolutional regression (conv_regression_blobs.py) and shared clause pools (mnist_conv.py --coalesced), and semantic segmentation on generated street scenes (segmentation_scenes.py).
• Worked notebooks in examples/notebooks/: Iris (rules you can read), MNIST (convolutional clauses and GPU throughput) and CIFAR-10 (booleanizing colour images), plus two that work through the Tsetlin-machine segmentation literature — the Convolutional Regression TM (ICMLT 2021) and CTM-UNet (ISTM 2025), the latter building a dense per-pixel Tsetlin segmentation head.

Citation

If you use torchtsetlin in research, please cite the Tsetlin machine papers whose algorithms you rely on (see the documentation) and this package:

@software{torchtsetlin,
  author = {Thambawita, Vajira},
  title  = {torchtsetlin: GPU-enabled, PyTorch-native Tsetlin machines},
  year   = {2026},
  url    = {https://github.com/vlbthambawita/torchtsetlin}
}

License

MIT — see LICENSE.