THRML is a JAX library for building and sampling probabilistic graphical models, with a focus on efficient block Gibbs sampling and energy-based models. Extropic is developing hardware to make sampling from certain classes of discrete PGMs massively more energy efficient; THRML provides GPU‑accelerated tools for block sampling on sparse, heterogeneous graphs, making it a natural place to prototype today and experiment with future Extropic hardware.
Features include:
- Blocked Gibbs sampling for PGMs
- Arbitrary PyTree node states
- Support for heterogeneous graphical models
- Discrete EBM utilities
- Enables early experimentation with future Extropic hardware
From a technical point of view, the internal structure compiles factor-based interactions to a compact "global" state representation, minimizing Python loops and maximizing array-level parallelism in JAX.
Installation
Requires Python 3.10+.
pip install thrml
or
uv pip install thrml
Documentation
Available at docs.thrml.ai.
Citing THRML
If you use THRML in your research, please cite us!
@misc{jelinčič2025efficientprobabilistichardwarearchitecture,
title={An efficient probabilistic hardware architecture for diffusion-like models},
author={Andraž Jelinčič and Owen Lockwood and Akhil Garlapati and Guillaume Verdon and Trevor McCourt},
year={2025},
eprint={2510.23972},
archivePrefix={arXiv},
primaryClass={cs.LG},
url={https://arxiv.org/abs/2510.23972},
}
Quick example
Sampling a small Ising chain with two-color block Gibbs:
import jax
import jax.numpy as jnp
from thrml import SpinNode, Block, SamplingSchedule, sample_states
from thrml.models import IsingEBM, IsingSamplingProgram, hinton_init
nodes = [SpinNode() for _ in range(5)]
edges = [(nodes[i], nodes[i+1]) for i in range(4)]
biases = jnp.zeros((5,))
weights = jnp.ones((4,)) * 0.5
beta = jnp.array(1.0)
model = IsingEBM(nodes, edges, biases, weights, beta)
free_blocks = [Block(nodes[::2]), Block(nodes[1::2])]
program = IsingSamplingProgram(model, free_blocks, clamped_blocks=[])
key = jax.random.key(0)
k_init, k_samp = jax.random.split(key, 2)
init_state = hinton_init(k_init, model, free_blocks, ())
schedule = SamplingSchedule(n_warmup=100, n_samples=1000, steps_per_sample=2)
samples = sample_states(k_samp, program, schedule, init_state, [], [Block(nodes)])
Sampling backends
Sampling uses JaxBackend by default. IsingSamplingProgram holds the logical model
in .model and its block layout in .gibbs_spec. Pass backend=your_backend to
sample_states, sample_with_observation, estimate_moments, or estimate_kl_grad
to use a backend supplied by an external package. Physical execution and calibration
belong to that package.
An initial state of None requests continuation of the backend's retained chain;
actual JAX sampling requires an explicit state. Use [] for an explicitly empty
state when there are no free blocks. Explicit physical initialization is optional
and backend-owned, and may require additional device operations.
For low-level JAX updates or access to .samplers and .per_block_*, convert an
Ising program with JaxBackend().compile(program). It uses a private JIT helper to
reuse the graph and layout work while applying the current model parameters on
every call. The first conversion compiles this helper; an enclosing JIT inlines it.
The result is the existing executable representation. Compile again after changing
the model parameters. Existing BlockSamplingProgram and FactorSamplingProgram
objects are returned unchanged.
See the concepts guide for portable observers,
backend batching, and training.
Sampling schedules are immutable because JAX uses them as compilation-cache keys.
Mutating a cached schedule can reuse a computation with the old record count.
Use dataclasses.replace(schedule, n_samples=100) to create an updated schedule.
More Examples
For a more extensive example, see the dtm-replication repository which runs on THRML.
