Getting started
THRML is a JAX library for building and sampling probabilistic graphical models, with a focus on efficient block Gibbs sampling and energy-based models. This page installs THRML and samples a first Ising chain.
Installation
THRML requires Python 3.10 or newer.
pip install thrmlOr with uv:
uv pip install thrmlFor a local checkout:
git clone https://github.com/extropic-ai/thrml
cd thrml
pip install -e .A first model
THRML samples graphical models by block Gibbs. Define the variables as nodes, wire their interactions into a model, divide the graph into blocks via graph-colouring so each block can update in parallel, then hand it all to a sampling program. Here is a small Ising chain end to end:
import jax
import jax.numpy as jnp
from thrml import SpinNode, Block, SamplingSchedule, sample_states
from thrml.models import IsingEBM, IsingSamplingProgram, hinton_init
# A 5-spin Ising chain: five nodes, four nearest-neighbour edges.
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)
# Two-colour the chain so each block updates independently under Gibbs.
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)])Each call to sample_states warms up the chain, then draws samples by alternating Gibbs updates over the two blocks. Because the even and odd spins never share an edge, every spin in a block resamples at once, which is exactly the structure Extropic's hardware is built to accelerate.
A Block is an ordered collection of nodes of the same type, and it is the unit block Gibbs updates together. Graph-colouring so that no two neighbours land in the same block is what lets a whole block resample in parallel.