thrml

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.