Hamon#
JAX-native thermal sampling for discrete energy-based models.
Hamon gives you GPU-accelerated block Gibbs sampling and non-reversible parallel tempering for Ising models, Boltzmann machines, and other discrete energy-based models — all in pure JAX.
import jax
from hamon import SpinNode, Block, SamplingSchedule, sample_states
from hamon.models import IsingEBM, IsingSamplingProgram, hinton_init
import jax.numpy as jnp
# Define a 5-spin Ising chain
nodes = [SpinNode() for _ in range(5)]
edges = [(nodes[i], nodes[i + 1]) for i in range(4)]
model = IsingEBM(nodes, edges, jnp.zeros(5), jnp.ones(4) * 0.5, jnp.array(1.0))
# Two-color block Gibbs
free_blocks = [Block(nodes[::2]), Block(nodes[1::2])]
program = IsingSamplingProgram(model, free_blocks, [])
key = jax.random.key(0)
k1, k2 = jax.random.split(key)
state = hinton_init(k1, model, free_blocks, ())
schedule = SamplingSchedule(n_warmup=100, n_samples=1000, steps_per_sample=2)
samples = sample_states(k2, program, schedule, state, [], [Block(nodes)])
Why "Hamon"?#
In Japanese swordsmithing, the hamon (刃文) is the visible wave along a katana's edge — a pattern created entirely by differential hardening. The smith coats the blade in clay, heats it to critical temperature, and quenches it. The edge cools fast into hard martensite; the spine cools slowly into tough pearlite. The hamon is the boundary between phases, born from a thermal process.
This library does the same thing computationally. It runs MCMC chains at different temperatures and exchanges information across the thermal gradient. Structure emerges at the boundary between mixing regimes — hot chains explore, cold chains resolve detail, and communication between them is what makes sampling work.
What Hamon provides#
Block Gibbs sampling on sparse, heterogeneous factor graphs with JAX-native compilation. Define your model as nodes, factors, and blocks; Hamon handles the index bookkeeping and padding for you.
Non-reversible parallel tempering (NRPT) with single-pass DEO swaps, adaptive schedule optimization, and round-trip diagnostics. Based on the theoretical framework of Syed et al. (2021).
Automatic chain count discovery via discover_chain_count, which estimates
the communication barrier Λ and recommends how many tempering chains you need.
Dynamic block management with influence-aware partitioning, per-temperature block sizing, and correlation-based re-blocking.
Origin#
Hamon began as a performance fork of Extropic AI's THRML library. It has since diverged into an independent project with its own algorithmic contributions. The original work is gratefully acknowledged under the Apache 2.0 license.