THRML

THRML 是一个基于 JAX 的库，用于构建和采样概率图模型，特别专注于高效的块吉布斯采样和基于能量的模型。Extropic 正在开发硬件，以大幅提高从特定类型的离散概率图模型中采样的能效；THRML 提供了针对稀疏、异构图的 GPU 加速块采样工具，使其成为当前进行原型设计以及未来与 Extropic 硬件协同实验的理想平台。

主要特性包括：

• 针对概率图模型的块吉布斯采样
• 任意 PyTree 节点状态
• 支持异构图模型
• 离散 EBM 工具集
• 便于早期探索未来 Extropic 硬件

从技术角度来看，其内部结构会将基于因子的交互编译为紧凑的“全局”状态表示，从而减少 Python 循环，并在 JAX 中最大化数组级别的并行性。

安装

需要 Python 3.10 或更高版本。

pip install thrml

或者

uv pip install thrml

文档

文档可在 docs.thrml.ai 查阅。

引用 THRML

如果您在研究中使用了 THRML，请引用我们！

@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}, 
}

快速示例

使用双色块吉布斯采样一个小伊辛链：

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)])

THRML 快速上手指南

THRML 是一个基于 JAX 的库，专为构建和采样概率图模型（PGM）而设计。它专注于高效的块吉布斯采样（Block Gibbs Sampling）和能量模型（EBM），特别适用于稀疏、异构图的 GPU 加速采样，也是探索未来 Extropic 硬件架构的理想原型工具。

环境准备

在开始之前，请确保您的开发环境满足以下要求：

• Python 版本：3.10 或更高版本
• 前置依赖：需已安装支持 GPU 加速的 JAX 环境（推荐）。若未安装，可先执行 pip install "jax[cuda]" 或访问 JAX 官网 获取安装指引。

安装步骤

您可以使用 pip 或 uv 进行安装。国内开发者如遇网络问题，建议配置国内镜像源（如清华源、阿里源）以加速下载。

使用 pip 安装：

pip install thrml

（可选）使用国内镜像源加速：

pip install thrml -i https://pypi.tuna.tsinghua.edu.cn/simple

使用 uv 安装：

uv pip install thrml

基本使用

以下示例演示了如何使用 THRML 对一个小型伊辛链（Ising chain）进行双色块吉布斯采样。

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, ())

# 配置采样计划：预热 100 步，采集 1000 个样本，每样本 2 步
schedule = SamplingSchedule(n_warmup=100, n_samples=1000, steps_per_sample=2)

# 执行采样
samples = sample_states(k_samp, program, schedule, init_state, [], [Block(nodes)])

更多详细文档请访问：docs.thrml.ai