THRML: Block Gibbs Sampling for Discrete PGMs in JAX

The problem THRML solves: sampling sparse, heterogeneous graphs without Python loops

Most probabilistic programming stacks are built around continuous variables and gradient-based inference. THRML starts from the opposite end. It targets discrete probabilistic graphical models where the natural inference move is Gibbs sampling, and where the graph is sparse and heterogeneous rather than a dense array of uniform factors. The README states the focus directly: efficient block Gibbs sampling and energy-based models.

The audience is narrow and identifiable. Extropic is developing hardware to make sampling from certain classes of discrete PGMs more energy efficient, and the README frames THRML as a place to prototype today and experiment with future Extropic hardware. So the primary user is a researcher who wants to write a discrete energy-based model now, run it on a GPU through JAX, and keep the model definition close enough to the hardware target that the same code stays meaningful later. The secondary user is anyone who has a sparse graph with non-uniform node types and finds that hand-written Gibbs loops in Python are the bottleneck.

The design answer to that bottleneck is stated in the README: factor-based interactions compile to a compact global state representation, which minimizes Python loops and maximizes array-level parallelism in JAX. That single sentence is the whole architectural thesis. Everything else in the API follows from it.

How the global state compilation works, and why arbitrary PyTree node states matter

The README lists arbitrary PyTree node states as a feature. That is not decoration. In a graph where one node holds a spin and another holds a categorical variable of different dimension, a naive implementation stores a list of heterogeneous objects and loops over it. Under JAX that loop is fatal: each iteration becomes a separate traced operation, and the sampler stops being a single compiled program.

THRML's answer is to compile the factor-based interactions into one compact global representation. The user still writes the model in terms of nodes, edges, biases and weights, but the sampler operates on the flattened state. Because node states can be arbitrary PyTrees, the flattening can carry structured per-node data rather than only scalars, while still presenting JAX with arrays it can vectorize over.

The Ising example in the README shows the shape of this. Nodes are SpinNode objects, edges are tuples of node pairs, and the model is constructed as IsingEBM(nodes, edges, biases, weights, beta). Sampling is then driven by two separate objects. A Block groups nodes that will be updated together, and the example splits a five-node chain into free_blocks = [Block(nodes[::2]), Block(nodes[1::2])], the two-color partition. A SamplingSchedule holds n_warmup, n_samples and steps_per_sample. The call to sample_states takes the program, the schedule, an initial state, and a list of blocks to record.

The two-color split is the mechanism, not an example detail. Blocked Gibbs only parallelizes if the nodes inside a block are conditionally independent given the rest, which is exactly what a bipartite partition of a chain gives you. On a general sparse graph the coloring problem is harder, and the README gives no partitioner. Choosing free_blocks is the user's job, and a bad partition either serializes the sampler or changes the stationary distribution.

Getting it running: install, imports, and the keys that actually matter

Installation is one command. The README requires Python 3.10 or newer and gives both forms:

pip install thrml

or

uv pip install thrml

There is no mention of CUDA wheels, a separate extras target, or a JAX version pin in the supplied material. Since THRML is a JAX library, the JAX install is presumably the user's responsibility, and the README does not say which jaxlib build is expected. That is worth checking before you debug anything else.

The quick example imports jax, jax.numpy as jnp, and from thrml the names SpinNode, Block, SamplingSchedule and sample_states, plus from thrml.models the names IsingEBM, IsingSamplingProgram and hinton_init. So the package is split into a core module and a models module, and the Ising machinery lives in the latter. If you are building a non-Ising model, expect to work against the core sampling API rather than the models helpers.

The keys that carry behavior are few. beta is a jnp.array scalar passed into IsingEBM, and it sets the inverse temperature of the energy model. free_blocks and clamped_blocks are passed to IsingSamplingProgram, and the example passes an empty list for clamped_blocks. SamplingSchedule(n_warmup=100, n_samples=1000, steps_per_sample=2) controls how many sweeps run before recording, how many samples come back, and how many Gibbs steps separate consecutive recorded samples. The final argument to sample_states is [Block(nodes)], the list of blocks whose states are collected. Setting steps_per_sample to 1 gives you correlated consecutive states; raising it thins the chain. The README does not state a default or a recommended value.

Randomness is handled with jax.random.key(0) and jax.random.split, and hinton_init takes its own key. That split matters: initialization and sampling consume separate streams, so reusing one key across both would silently correlate them.

What the API does not cover: continuous variables, inference, and graph partitioning

The honest limitation is scope. Every concrete example in the README is an Ising model over SpinNode objects with scalar biases and weights. The library advertises support for heterogeneous graphical models and arbitrary PyTree node states, but the supplied material contains no worked example of a heterogeneous model, so the practical ergonomics of that path cannot be judged from the README alone. Treat the heterogeneous claim as documented but unillustrated.

There is no mention of continuous latent variables, no variational inference, no Hamiltonian Monte Carlo, no autoguide, and no model specification language. If your problem is a hierarchical regression with continuous parameters, THRML is the wrong tool and no amount of configuration will change that. It samples discrete state spaces.

The second limitation is the partition. Blocked Gibbs requires you to supply the blocks, and the README's example uses the trivial two-coloring of a chain. On a sparse graph with odd cycles, a two-coloring may not exist, and the README offers no guidance on what to do. You can fall back to single-node blocks, which is correct but gives up the parallelism the library exists to provide. This is the point where a prototype either scales or quietly becomes a slow Python loop wearing a JAX costume.

The third is version risk. The release list shows v0.1.3 in October 2025 and v0.1.4 in August 2026, with v0.1.3 labeled Initial Release. That is a short history at the 0.1 line, and the README carries no compatibility or deprecation policy. Pin the version you test against.

Where THRML sits against PyMC and NumPyro

The closest comparison is NumPyro, which is also JAX-based and also compiles sampling into XLA. The difference is what gets compiled. NumPyro takes a generative model written in its primitives and runs NUTS or a discrete Gibbs sampler over the resulting log-density, with the graph implied by the model code. THRML takes an explicit graph of nodes and edges, an explicit energy function in the Ising case, and an explicit block partition, then runs blocked Gibbs over the compiled global state. NumPyro hides the graph; THRML makes you declare it.

That trade is deliberate. Declaring the graph is what lets THRML compile factor interactions into one flat state and parallelize a block update across many nodes at once. NumPyro's generic sampler cannot assume your graph has a useful coloring, so it cannot make that move. If your model is discrete and sparse, THRML's explicit graph is the price of the parallelism. If your model is continuous, or if you want NUTS and posterior predictive checks out of the box, NumPyro is the more direct route and THRML has nothing to offer you.

PyMC is a further step away, being a Python-level modeling language with a backend that has historically been gradient-based. The comparison is only useful as a boundary marker: THRML is not a modeling language, it is a sampling runtime with a small model layer on top. You bring the graph.

One more data point from the README: the dtm-replication repository is cited as a more extensive example that runs on THRML. That is the only external usage the supplied material points to, and it is the right place to look before writing your own model from scratch.

Licence, maintenance, and what upgrading actually costs

THRML is Apache-2.0. That permits commercial use, modification and redistribution provided the licence and notices are preserved, and it includes an explicit patent grant. It is a permissive licence, not a copyleft one, so it does not force you to open your own model code. This is a description of the licence text, not legal advice; if the patent grant or the notice requirements matter to your organization, have counsel read the file.

Maintenance signals in the supplied material are limited to the release list and the last push timestamp. Two releases exist: v0.1.3, labeled Initial Release, in October 2025, and v0.1.4 in August 2026. The gap between them is roughly nine months, and the repository was pushed in September 2026. There is no changelog in the material, so what changed between v0.1.3 and v0.1.4 is not knowable from what is here. The documentation site at docs.thrml.ai is the place that would answer it.

Upgrade cost is hard to estimate from two releases at the 0.1 line. The practical exposure is that the API surface is small enough that a rename in thrml.models would touch every model file you have written, and there is no stated deprecation policy. The mitigation is cheap: pin thrml in your dependency file, and check the docs site for a migration note before moving between minor versions. If you are running this as part of a paper replication, pin the exact version the replication used.

The citation block in the README points to an arXiv paper on a probabilistic hardware architecture for diffusion-like models, with Jelinčič, Lockwood, Garlapati, Verdon and McCourt as authors. If you are evaluating whether the sampling semantics match your intended model, that paper is the reference the maintainers themselves direct you to, and it is more likely to explain the block structure than the README is.

Editorial conclusion

Adopt THRML if you are prototyping discrete energy-based models and want block Gibbs sampling to run inside JAX transformations, and if you accept that the API is organized around Ising-style factors, SpinNode and Block objects, and a SamplingSchedule rather than a general model specification language. Do not adopt it if you need continuous latent variables, variational inference, or NUTS; the README documents none of those, and the library's stated purpose is discrete sampling. Before committing, verify three things: that your Python is 3.10 or newer as the installation section requires, that your graph can be partitioned into the free_blocks you pass to IsingSamplingProgram, and that the pinned version v0.1.4 behaves as the docs.thrml.ai examples describe on your hardware, since the jump from v0.1.3 to v0.1.4 is the only release history available and the API surface is still small.

Official sources

1. extropic-ai/thrml on GitHub
2. License: Apache-2.0
3. Project website
4. README
5. Releases