THRML: Thermodynamic Hypergraphical Model Library

THRML (Thermodynamic Hypergraphical Model Library) is a JAX-based Python library designed for building and efficiently sampling probabilistic graphical models (PGMs). Developed by Extropic, it focuses on advanced techniques such as block Gibbs sampling and the use of energy-based models (EBMs).

Purpose and Vision

The core purpose of THRML is to provide GPU-accelerated tools for simulating and experimenting with sparse, heterogeneous graphical models. It serves as a crucial prototyping platform for Extropic's future hardware, which aims to achieve significantly more energy-efficient sampling from specific classes of discrete PGMs. By offering robust tools for today's hardware, THRML enables researchers and developers to explore and validate concepts that will eventually leverage Extropic's specialized computational infrastructure.

Key Features

THRML is equipped with a suite of features tailored for modern PGM research and development:

• Blocked Gibbs sampling for PGMs: Implements an optimized approach to Gibbs sampling where groups of variables (blocks) are updated simultaneously, improving convergence and computational efficiency.
• Arbitrary PyTree node states: Supports flexible and complex node state representations using JAX's PyTree structure, allowing for diverse data types and nested structures within graph nodes.
• Support for heterogeneous graphical models: Capable of handling graphs where nodes can be of different types (e.g., binary spins, categorical variables) and interact in varied ways.
• Discrete EBM utilities: Provides specific tools and functionalities for constructing and manipulating discrete energy-based models.
• Enables early experimentation with future Extropic hardware: Acts as a sandbox for developing and testing algorithms compatible with the unique capabilities of upcoming Extropic hardware.

Technical Architecture

Under the hood, THRML's architecture is engineered for performance within the JAX ecosystem. It translates factor-based interactions into a compact "global" state representation. This design choice is fundamental to minimizing traditional Python loops, thereby maximizing array-level parallelism inherent to JAX. This optimized structure ensures that operations on large, complex graphs can be executed with high efficiency on accelerators like GPUs.