# Table of Contents - [Developer Documentation - THRML](#developer-documentation-thrml) - [THRML](#thrml) - [Getting Started with THRML - THRML](#getting-started-with-thrml-thrml) - [Graphical Model Components - THRML](#graphical-model-components-thrml) - [Block Management - THRML](#block-management-thrml) - [Interaction Groups - THRML](#interaction-groups-thrml) - [Factors - THRML](#factors-thrml) - [Block Sampling - THRML](#block-sampling-thrml) - [Ising Models - THRML](#ising-models-thrml) - [Sampling Observers - THRML](#sampling-observers-thrml) - [Discrete Energy-Based Models - THRML](#discrete-energy-based-models-thrml) - [Energy-Based Models - THRML](#energy-based-models-thrml) - [Conditional Samplers - THRML](#conditional-samplers-thrml) - [Spin Models in THRML - THRML](#spin-models-in-thrml-thrml) - [All of THRML - THRML](#all-of-thrml-thrml) - [THRML](#thrml) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) - [Just a moment...](#just-a-moment-) --- # Developer Documentation - THRML [Skip to content](https://docs.thrml.ai/en/latest/architecture/#developer-documentation) Developer Documentation[¤](https://docs.thrml.ai/en/latest/architecture/#developer-documentation "Permanent link") =================================================================================================================== What is `THRML`?[¤](https://docs.thrml.ai/en/latest/architecture/#what-is-thrml "Permanent link") -------------------------------------------------------------------------------------------------- As was discussed in previous documents, `THRML` is a [JAX](https://docs.jax.dev/en/latest/) ‑based Python package for efficient [block Gibbs sampling](https://proceedings.mlr.press/v15/gonzalez11a/gonzalez11a.pdf) of graphical models at scale. `THRML` provides the tools to do block Gibbs sampling on any graphical model, and provides the tooling already for models such as [Boltzmann Machines](https://www.cs.toronto.edu/~fritz/absps/cogscibm.pdf) . How does `THRML` work?[¤](https://docs.thrml.ai/en/latest/architecture/#how-does-thrml-work "Permanent link") -------------------------------------------------------------------------------------------------------------- From a user perspective, there are three main components of `THRML` that they will interact with: blocks, factors, and programs. For detailed usage examples, see the example notebooks. Blocks are fundamental to `THRML` since it implements block sampling. A `Block` is a collection of nodes of the same type with implicit ordering. Factors and their associated conditionals are the backbone of sampling. Factors derive their name from factor graphs, and organize interactions between variables into a [bipartite graph of factors and variables](https://ocw.mit.edu/courses/6-438-algorithms-for-inference-fall-2014/3e3e9934d12e3537b4e9b46b53cd5bf1_MIT6_438F14_Lec4.pdf) . Factors synthesize collections of interactions via `InteractionGroups` and must implement a `to_interaction_groups()` method. Below is the hierarchy of interactions and samplers provided for clarity. Programs are the key orchestrating data structures. `BlockSamplingProgram` handles the mapping and bookkeeping for padded block Gibbs sampling, managing global state representations efficiently for JAX. `FactorSamplingProgram` is a convenient wrapper that converts factors to interaction groups. These programs coordinate free/clamped blocks, samplers, and interactions to execute the sampling algorithm. From a developer perspective, the core approach to `THRML` is to represent as much as possible as contiguous arrays/pytrees, operate on these structures, then map to and from them for the user. Internally, this is often referred to as the "global" state (in opposition to the "block" state). This can be seen as a similar approach to data driven design (via SoA) and is similar to other JAX graphical model packages (e.g. [PGMax](https://github.com/google-deepmind/PGMax) ). Taking PGMax as an example, an important distinction is that `THRML` supports pytree states and heterogeneous states. There is more than one way to approach this heterogeneity and `THRML` takes an approach that relies on splitting the nodes according to the pytrees and organizing a global state as a list of these pytrees (which are then stacked if there are multiple blocks that share a given pytree). Thus, the global state is a list of these pytree structures. Since JAX is optimized for efficient array/pytree operation we want to do as much as possible in that form, so we define a standard representation and order for this global structure (which itself doesn't really matter much, it just matters that we know how to get to and from this order) in this array format (in which all pytree structs of the same type get stacked together), then map indices there and back to other representations. The management of these indices and mapping is constructed/held by the program. Since JAX does not support ragged arrays, every block must be the same size (in the array leaves). In order to solve this problem (since blocks in the graph may be different sizes), `THRML` constructs the global representation by stacking the blocks (of the same pytree type) and pad them out as needed. There exists a tradeoff between padding out blocks which can add runtime overhead (from unnecessary computation) and other approaches, such as looping over blocks which could pay (a likely untenable) compile time cost instead. Everything else that exists in `THRML` exists to provide convenience for creating and working with a program. With a focused core on block index management and padding, this allows for a lightweight and hackable code base (with only 1,000 LoC). What are the limitations of `THRML`?[¤](https://docs.thrml.ai/en/latest/architecture/#what-are-the-limitations-of-thrml "Permanent link") ------------------------------------------------------------------------------------------------------------------------------------------ While `THRML` is fast and efficient, users new to sampling may expect a panacea where none can exist. First and foremost, it is important to note that sampling is a very difficult problem. To generate samples from a distribution in high dimensional space can take (prohibitively) many steps even if we parallelize proposals. `THRML` is also very focused on Gibbs sampling, as Extropic seeks to provide hardware that accelerates this algorithm, but for general sampling it is unknown when Gibbs sampling as an MCMC method is substantially [faster](https://arxiv.org/abs/2007.08200) or [slower](https://arxiv.org/abs/1605.00139) than other MCMC methods and thus specific problems may require specific tools. As a pedagogical example, consider a two node Ising model with a single edge. If J\=−∞,h\=0, Gibbs sampling will never mix between the ground states {-1, -1}, {1, 1} since it will never flip once it reaches one of these states (but an approach such as Uniform MH would be able to converge quickly). `THRML` Overviews[¤](https://docs.thrml.ai/en/latest/architecture/#thrml-overviews "Permanent link") ----------------------------------------------------------------------------------------------------- ![A diagram which shows the flow of different components into the FactorSamplingProgram](https://docs.thrml.ai/en/latest/flow.png) #### Factors:[¤](https://docs.thrml.ai/en/latest/architecture/#factors "Permanent link") * `AbstractFactor` * `WeightedFactor`: Parameterized by weights * `EBMFactor`: defines energy functions for Energy-Based Models * `DiscreteEBMFactor`: EBMs with discrete states (spin and categorical) * `SquareDiscreteEBMFactor`: Optimized for square interaction tensors * `SpinEBMFactor`: Spin-only interactions ({-1, 1} variables) * `SquareCategoricalEBMFactor`: Square categorical interactions * `CategoricalEBMFactor`: Categorical-only interactions #### Samplers:[¤](https://docs.thrml.ai/en/latest/architecture/#samplers "Permanent link") * `AbstractConditionalSampler` * `AbstractParametricConditionalSampler` * `BernoulliConditional`: Spin-valued Bernoulli sampling * `SpinGibbsConditional`: Gibbs updates for spin variables in EBMs * `SoftmaxConditional`: Categorical softmax sampling * `CategoricalGibbsConditional`: Gibbs updates for categorical variables in EBMs --- # THRML [Skip to content](https://docs.thrml.ai/en/latest/#thermodynamic-hypergraphical-model-library-thrml) ![THRML Logo](https://docs.thrml.ai/en/latest/_static/logo/logo.svg) **Thermodynamic HypergRaphical Model Library (THRML)**[¤](https://docs.thrml.ai/en/latest/#thermodynamic-hypergraphical-model-library-thrml "Permanent link") ============================================================================================================================================================== * * * `THRML` is a JAX library for building and sampling probabilistic graphical models, with a focus on efficient block Gibbs sampling and energy-based models. Extropic is developing hardware to make sampling from certain classes of discrete PGMs massively more energy‑efficient; `THRML` provides GPU‑accelerated tools for block sampling on sparse, heterogeneous graphs, making it a natural place to prototype today and experiment with future Extropic hardware. Features include: * Block Gibbs sampling for PGMs * Arbitrary PyTree node states * Support for heterogeneous graphical models * Discrete EBM utilities (Ising/RBM‑like) * Enables early experimentation with future Extropic hardware Installation[¤](https://docs.thrml.ai/en/latest/#installation "Permanent link") -------------------------------------------------------------------------------- Requires >=python 3.10 `pip install thrml` or `uv pip install thrml` For installing from the source: `git clone https://github.com/extropic-ai/thrml cd thrml pip install -e .` or `git clone https://github.com/extropic-ai/thrml cd thrml uv pip install -e .` Quick example[¤](https://docs.thrml.ai/en/latest/#quick-example "Permanent link") ---------------------------------------------------------------------------------- Sampling a small Ising chain with two‑color block Gibbs: `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)])` --- # Getting Started with THRML - THRML [Skip to content](https://docs.thrml.ai/en/latest/examples/00_probabilistic_computing/#getting-started-with-thrml) Getting Started with THRML[¤](https://docs.thrml.ai/en/latest/examples/00_probabilistic_computing/#getting-started-with-thrml "Permanent link") ================================================================================================================================================ Extropic hardware will become increasingly capable over the next few years. If everything goes to plan, our first ~million variable chips will come online in 2026 and be integrated into large many-chip systems by 2027. These systems will be capable of complex probabilistic computations and could offer a substantial edge to anyone who knows how to wield them. Given these (relatively) short timelines to large-scale commercial viability, businesses must start working today to adapt our tech to their use cases. To enable this, we built a software library that simulates Extropic hardware devices on traditional machine learning accelerators like GPUs. Users can leverage this library to explore how future Extropic systems could be used to accelerate workloads they care about. Our simulation library, THRML, is a GPU simulator of the probabilistic sampling programs that run natively on Extropic hardware. It is built on top of JAX and is massively scalable, enabling users to simulate arbitrarily large probabilistic computing systems given sufficient GPU resources. What's a probabilistic computer?[¤](https://docs.thrml.ai/en/latest/examples/00_probabilistic_computing/#whats-a-probabilistic-computer "Permanent link") ---------------------------------------------------------------------------------------------------------------------------------------------------------- Probabilistic computers leverage arrays of massively parallel random number generators to sample from distributions that are relevant to solving practical problems. By leveraging physical effects to generate random numbers and emphasising local communication, probabilistic computers can process information extremely efficiently. These sampling-centric computers make direct contact with contemporary machine learning via _Energy Based Models_ (EBMs). Much like popular diffusion models and autoregressive language models, EBMs attempt to learn the probability distribution that underlies some observations of the world. However, EBMs are unique in that they try to directly learn the shape of this distribution instead of attempting to represent it via the composition of many much simpler distributions. This direct fitting of data distributions leads to training and inference pipelines that are extremely sampling-heavy. While these sampling-heavy workloads can challenge traditional machine learning accelerators, they present a massive opportunity for probabilistic computers. In particular, symmetries of certain families of EBMs can be leveraged to map EBM sampling problems directly onto probabilistic computing hardware, unlocking ultra-efficient EBM training and inference. We leveraged these symmetries, along with other hardware and algorithm innovations, to demonstrate that probabilistic computers can be many orders of magnitude more energy-efficient than GPUs [in our recent paper](https://arxiv.org/abs/2510.23972) . How do Extropic's probabilistic computers work?[¤](https://docs.thrml.ai/en/latest/examples/00_probabilistic_computing/#how-do-extropics-probabilistic-computers-work "Permanent link") ---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- To truly understand how Extropic's probabilistic computers operate and their utility in machine learning, we must examine the underlying mathematics of EBMs and the algorithms that enable probabilistic computers to sample from them. EBMs attempt to model data distributions by learning a parameterized _Energy Function_ E(x,θ) that defines the shape of the model's probability distribution, P(X\=x)∝e−E(x,θ) Here, X is a vector of random variables that represents the data you want to model, which could be text, images, etc. EBMs are fit by tweaking the model's parameters θ to assign low values of energy to values of x that show up in the dataset and large values of energy to values of x that don't. For EBMs, both inference and training involve sampling from this potentially very complex probability distribution, which is very expensive to do on a CPU or GPU. At inference time, we aim to generate new data that resembles the dataset, which requires drawing samples from the learned distribution. During training, to estimate gradients of the typical EBM training objective with respect to the parameters θ, one needs to compute estimators of certain averages over the EBM's distribution, which means lots of sampling. If a probabilistic computer can make this sampling very efficient, it could dramatically improve the real-world performance of EBMs relative to other types of machine learning models. Extropic's probabilistic computers can efficiently sample from _factorized_ EBMs that contain only _local_ interactions. An EBM is factorized if its energy function splits up into a sum over independent terms, E(x)\=∑(ψ,x1,…,xN)∈Sψ(x1,…,xN) where S is the set of all the factors involved in the EBM. Each member of S consists of a factor energy function ψ that acts on a subset of the model's variables x1,…,xN. A factor is _local_ if it only involves variables that are somehow "close" to each other, which in the context of a probabilistic computer means that they are embodied by circuitry that lives on nearby parts of a chip. Extropic's probabilistic computers leverage the Gibbs sampling algorithm to efficiently sample from these special EBMs. Gibbs sampling is a procedure in which samples are drawn from the EBM's distribution by iteratively updating blocks of non-interacting variables according to their conditional distributions. P(Xi\=xi|Xnb(i)\=x′)∝e−Ei(xi,x′) where Xnb(i) is the set of random variables that Xi interacts with and Ei contains contributions from the set of factors that involve the state of Xi, which we denote as Si, Ei(xi,x′)\=∑(ψ,x1,…,xK)∈Siψ(xi,x1,…,xK) Extropic's probabilistic computers implement this block Gibbs sampling procedure at the hardware level to dramatically reduce the energy cost of EBM training and inference. A concrete example[¤](https://docs.thrml.ai/en/latest/examples/00_probabilistic_computing/#a-concrete-example "Permanent link") -------------------------------------------------------------------------------------------------------------------------------- To make all of this more concrete, here we will implement a simulation of a simple probabilistic computer sampling from an EBM using THRML. We will implement sampling for a _Potts Model_, which is a type of EBM that was first developed to study various phenomena in solid-state physics. Potts models have energy functions like, E(x)\=−β(∑iWi(1)\[xi\]+∑(i,j)∈SWi,j(2)\[xi,xj\]) Here, xi is an integer representing a possible state of a categorical random variable. The vector Wi(1) induces a bias on the ith variable by adding or subtracting energy based on the value of xi. The matrix Wi,j(2) generates an interaction between the ith and jth variables by contributing a term to the energy function that depends on both xi and xj. Let's use THRML to run block Gibbs sampling on a Potts model. First, some imports that will be useful, `import jax import jax.numpy as jnp import matplotlib.pyplot as plt import networkx as nx` `from thrml.block_management import Block from thrml.block_sampling import BlockGibbsSpec, sample_states, SamplingSchedule from thrml.pgm import CategoricalNode from thrml.models.discrete_ebm import CategoricalEBMFactor, CategoricalGibbsConditional from thrml.factor import FactorSamplingProgram` The Potts model energy function naturally suggests a graphical interpretation. Namely, we can assign each random variable in the problem to a node, and assign biases to each node to implement the W(1). We can connect variables with edges to represent the interactions W(2). As such, let's define a simple graph to use in our problem, `side_length = 20 # in this simple example we will just use a basic grid, although THRML is capable of dealing with arbitrary graph topologies G = nx.grid_graph(dim=(side_length, side_length), periodic=False) # label the nodes with something THRML recognizes for convenience coord_to_node = {coord: CategoricalNode() for coord in G.nodes} nx.relabel_nodes(G, coord_to_node, copy=False) for coord, node in coord_to_node.items(): G.nodes[node]["coords"] = coord # write down the color groups for later bicol = nx.bipartite.color(G) color0 = [n for n, c in bicol.items() if c == 0] color1 = [n for n, c in bicol.items() if c == 1] # write down the edges in a different format for later u, v = map(list, zip(*G.edges())) # plot the graph pos = {n: G.nodes[n]["coords"][:2] for n in G.nodes} colors = ["black", "orange"] node_colors = [colors[bicol[n]] for n in G.nodes] fig, axs = plt.subplots() nx.draw( G, pos=pos, ax=axs, node_size=50.0, node_color=node_colors, edgecolors="k", with_labels=False, )` ![img](https://docs.thrml.ai/en/latest/examples/_data/87bad8ce9f5f40648dfa6a9942896e48.png) The graph we just drew can be interpreted as a high-level schematic for a probabilistic computer. Each node is associated with specialized circuitry for performing the Gibbs sampling conditional update, which requires state information from neighboring nodes that is communicated along the edges (which could represent physical wires). Nodes of the same color can be updated in parallel, which means that a single iteration of Gibbs sampling can be performed by first updating all of the orange nodes simultaneously, followed by all of the blue nodes. Now that we have our graph, we can define the interactions that make up the energy function we described earlier. THRML comes with a canned implementation of Potts model style interactions in `thrml.models` that we can leverage here. We will consider a simple case that is commonly studied in physics with no biases and identity coupling matrices. `# how many categories to use for each variable n_cats = 5 # temperature parameter beta = 1.0 # implements W^{2} for each edge # in this case we are just using an identity matrix, but this could be anything id_mat = jnp.eye(n_cats) weights = beta * jnp.broadcast_to(jnp.expand_dims(id_mat, 0), (len(u), *id_mat.shape)) coupling_interaction = CategoricalEBMFactor([Block(u), Block(v)], weights) interactions = [coupling_interaction]` We can use these interactions to build a sampling program! THRML exposes tools that we can use to very easily run the Gibbs sampling algorithm for our problem. For our Potts model, we will need to update our state using a conditional distribution with the energy function, Ei(xi,x′)\=−β(Wi(1)\[xi\]+∑j∈SiWi,j(2)\[xi,xj\]) This corresponds to sampling from a _softmax_ distribution. THRML comes pre-packaged with this conditional, which we will leverage to perform our sampling: `# first, we have to specify a division of the graph into blocks that will be updated in parallel during gibbs sampling # During gibbs sampling, we are only allowed to update nodes in parallel if they are not neighbours # Mathematically, this means we should choose our sampling blocks based on the "minimum coloring" of the graph # we already computed this earlier # a Block of nodes is simply a sequence of nodes that are all the same type # we only have one type of node here, so not important to understand yet blocks = [Block(color0), Block(color1)] # our grouping of the graph into blocks spec = BlockGibbsSpec(blocks, []) # we have to define how each node in our blocks should be updated during each iteration of Gibbs sampling # THRML comes with a sampler that will do this for the vanilla potts model we are using here, so lets use that sampler = CategoricalGibbsConditional(n_cats) # now we can make a sampling program, which combines our grouping with the interactions we defined earlier prog = FactorSamplingProgram( spec, # our block decomposition of the graph [sampler for _ in spec.free_blocks], # how to update the nodes in each block every iteration of Gibbs sampling interactions, # the interactions present in our model [], )` That's everything! Now we can simply run our sampling program and observe the results, `# rng seed key = jax.random.key(4242) # everything in THRML is completely compatible with standard jax functionality like jit, vmap, etc. # here we will use vmap to run a bunch of parallel instances of Gibbs sampling n_batches = 100 # we need to initialize our Gibbs sampling instances init_state = [] for block in spec.free_blocks: key, subkey = jax.random.split(key, 2) init_state.append( jax.random.randint(subkey, (n_batches, len(block.nodes)), minval=0, maxval=n_cats, dtype=jnp.uint8) ) # how we should schedule our sampling schedule = SamplingSchedule( # how many iterations to do before drawing the first sample n_warmup=0, # how many samples to draw in total n_samples=100, # how many steps to take between samples steps_per_sample=5, ) keys = jax.random.split(key, n_batches) # now run sampling samples = jax.jit(jax.vmap(lambda i, k: sample_states(k, prog, schedule, i, [], [Block(G.nodes)])))(init_state, keys)` If we look at the results of our sample, we are able to observe domain formation, which is a very famous property of Potts models. Domain formation happens because our diagonal weight matrix encourages neighbouring variables to match,, leading to the formation of _domains_ where groups of neighbouring variables are all in the same state. `to_plot = [0, 7, 21] fig, axs = plt.subplots(nrows=1, ncols=len(to_plot)) for i, num in enumerate(to_plot): axs[i].imshow(samples[0][num, -1, :].reshape((side_length, side_length)))` ![img](https://docs.thrml.ai/en/latest/examples/_data/77f297f8cca24d40a9ad8e63374d2bb4.png) Despite its apparent simplicity, this example gets at the heart of why EBMs can be used as powerful machine learning primitives. Even though our model only involves simple interactions between neighbouring variables, it produces complex emergent long-range correlations that are difficult to predict and understand. By learning the weights of a model like this (instead of just setting them to be the identity matrix), we can take advantage of this capacity for complexity to model complex real-world phenomena, and do so using very little energy by leveraging probabilistic computing hardware. We've only scratched the surface of what can be done with THRML in this simple example. While here we mostly leaned on canned functionality built into THRML, in reality, THRML was designed from the ground up to make it easier for users to implement their own entirely novel probabilistic graphical model block-sampling routines. We designed it this way because probabilistic computing is a rapidly changing field. So it's essential to have flexible tools available to you that minimize exploratory friction. The other examples in this repo will begin to explore some of the more advanced functionality THRML offers. If you are interested in using THRML for your own research, we encourage you to check it out. Suppose you want to learn more about how hardware-compatible EBMs may be applied to real machine learning problems. In that case, you should check out [our paper](http://arxiv.org/abs/2510.23972) and an [implementation of it in `THRML`](https://github.com/pschilliOrange/dtm-replication) . Our primary hope for this library is that it empowers its users to build on our work, developing progressively more complex machine learning systems that increasingly leverage ultra-efficient probabilistic computing. --- # Graphical Model Components - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/pgm/#graphical-model-components) Graphical Model Components[¤](https://docs.thrml.ai/en/latest/api/pgm/#graphical-model-components "Permanent link") ==================================================================================================================== #### `thrml.AbstractNode` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/pgm.py#L58-70) [¤](https://docs.thrml.ai/en/latest/api/pgm/#thrml.AbstractNode "Permanent link") A node in a PGM. Every node used in a PGM must inherit from this class. When compiling a program, each node is assigned a shape and datatype that are used to organize the state of the sampling program in a jax-friendly way. #### `thrml.SpinNode([thrml.AbstractNode](https://docs.thrml.ai/en/latest/api/pgm/#thrml.AbstractNode) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/pgm.py#L73-76) [¤](https://docs.thrml.ai/en/latest/api/pgm/#thrml.SpinNode "Permanent link") A node that represents a random variable that takes on a state in {-1, 1}. #### `thrml.CategoricalNode([thrml.AbstractNode](https://docs.thrml.ai/en/latest/api/pgm/#thrml.AbstractNode) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/pgm.py#L79-83) [¤](https://docs.thrml.ai/en/latest/api/pgm/#thrml.CategoricalNode "Permanent link") A node that represents a random variable that may take on any one of K possible discrete states, represented by a positive integer in (0, K\]. --- # Block Management - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/block_management/#block-management) Block Management[¤](https://docs.thrml.ai/en/latest/api/block_management/#block-management "Permanent link") ============================================================================================================= #### `thrml.Block` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L30-81) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.Block "Permanent link") A Block is the basic unit through which Gibbs sampling can operate. Each block represents a collection of nodes that can efficiently be sampled simultaneously in a JAX-friendly SIMD manner. In THRML, this means that the nodes must all be of the same type. **Attributes:** * `nodes`: the tuple of nodes that this block contains #### `thrml.BlockSpec` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L88-207) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec "Permanent link") This contains the necessary mappings for logging indices of states and node types. This helps convert between block states and global states. A block state is a list of pytrees, where each pytree leaf has shape\[0\] = number of nodes in the block. The length of the block state is the number of blocks. The global state is a flattened version of this. Each pytree type is combined (regardless of which block they are in), to make a list of pytrees where each leaf shape\[0\] is the total number of nodes of that pytree shape. As an example, imagine an Ising model, every node is the same pytree (just a scalar array), as such the block state is a list of arrays where each array is the state of the block and the global state would be a length-1 list that contains an array of shape (total\_nodes,). Why is this global/block representation necessary? The answer is that the global representation is preferred for operating over in many JAX cases, but requires careful indexing (to know where in this long array each block resides) and thus the block representation is more natural/easy to use for many users. Why is the global state easier to work with? Well consider sampling, in order to sample a block (or even just a node) we need to collect all the states of the neighboring nodes. If we only had the block state we would have to loop over the block state and collect from each block the neighbors, we would then pass this to the sampler. The sampler would then have to know the type of each block (to know what to do with the states) then for loop over the blocks in order to collect each. This (programmatically) is fine, but results in additional for loops that slow down JAX, compared to gathering indexes from a single array. **Attributes:** * `blocks`: the list of blocks this spec contains * `all_block_sds`: a SD is a single `_PyTreeStruct`. Each node/block has only one SD associated with it, but each node can have neighbors of many types. This is the SD of each block (in the same order as blocks, this internal ordering is quite important for bookkeeping). This list is just the list of SDs for each block (and thus has length = len(blocks)). * `global_sd_order`: the list of SDs, providing a SoT for the global ordering * `sd_index_map`: a dictionary mapping the SD to an integer in the `global_sd_order`. This is like calling `.index` on it. * `node_global_location_map`: a dictionary mapping a given node to a tuple. That tuple contains the global index (i.e. which element in the global list it is in) and the relative position in that pytree. That is to say, you can get the state of the node via `map(x[tuple[1]], global_repr[tuple[0]])` * `block_to_global_slice_spec`: a list over unique SDs (so length global\_sd\_order), where each list inside this is the list over blocks which contain that pytree. E.g. \[\[0, 1\], \[2\]\] indicates that blocks\[0\] and blocks\[1\] are both of pytree SD 0. * `node_shape_dtypes`: a dictionary mapping node types to hashable `_PyTreeStruct` * `node_shape_struct`: a dictionary mapping node types to pytrees of JAX-shaped dtype structs (just for user access, since the keys aren't hashable that creates issues for JAX in other areas.) ##### `__init__(blocks: list[thrml.Block], node_shape_dtypes: Mapping[type[thrml.AbstractNode], PyTree[jax._src.core.ShapeDtypeStruct]])` [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec.__init__ "Permanent link") Create a BlockSpec from blocks. Based on the information passed in via node\_shape\_dtypes, determine the minimal global state that can be used to represent the blocks. **Arguments:** * `blocks`: the list of `Block`s that this specification operates on * `node_shape_dtypes`: the mapping of node types to their structures. This should be a pytree of `jax.ShapeDtypeStruct`s. #### `thrml.block_state_to_global(block_state: list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']], spec: thrml.BlockSpec) -> list[PyTree[Shaped[Array, 'nodes_global ?*state'], '_GlobalState']]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L221-258) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.block_state_to_global "Permanent link") Convert block-local state to the global stacked representation. The block representation is a list where `block_state[i]` contains the state of `spec.blocks[i]` and every node occupies index 0 of its leaf. The global representation is a shorter list (one entry per distinct PyTree structure) in which all blocks with the same structure are concatenated along their node axis. **Arguments:** * `block_state`: State organised per block, same length as `spec.blocks`. * `spec`: The [`thrml.BlockSpec`](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec) that defines the mapping. **Returns:** A list whose length equals `len(spec.global_sd_order)`—the stacked global state. #### `thrml.get_node_locations(nodes: thrml.Block, spec: thrml.BlockSpec) -> tuple[int, Int[Array, 'nodes']]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L261-290) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.get_node_locations "Permanent link") Locate a contiguous set of nodes inside the global state. **Arguments:** * `nodes`: A [`thrml.Block`](https://docs.thrml.ai/en/latest/api/block_management/#thrml.Block) whose nodes you want locations for. * `spec`: The [`thrml.BlockSpec`](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec) generated from the same graph. **Returns:** Tuple `(sd_index, positions)` where * _sd\_index_ is the position inside the global list returned by [`thrml.block_state_to_global`](https://docs.thrml.ai/en/latest/api/block_management/#thrml.block_state_to_global) , and * _positions_ is a 1D array with the indices each node occupies inside that particular PyTree. #### `thrml.from_global_state(global_state: list[PyTree[Shaped[Array, 'nodes_global ?*state'], '_GlobalState']], spec_from: thrml.BlockSpec, blocks_to_extract: list[thrml.Block]) -> list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L293-324) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.from_global_state "Permanent link") Extract the states for a subset of blocks from a global state. **Arguments:** * `global_state`: A state produced by \[`thrml.block_state_to_global(spec_from)`\]\[\]. * `spec_from`: The [`thrml.BlockSpec`](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec) associated with _global\_state_. * `blocks_to_extract`: The blocks whose node states should be returned. **Returns:** A list with one element per _blocks\_to\_extract_—each element is a PyTree with exactly `len(block)` nodes in its leading dimension. #### `thrml.make_empty_block_state(blocks: list[thrml.Block], node_shape_dtypes: Mapping[type[thrml.AbstractNode], PyTree[jax._src.core.ShapeDtypeStruct]], batch_shape: tuple | None = None) -> list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L327-361) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.make_empty_block_state "Permanent link") Allocate a zero-initialised block state. **Arguments:** * `blocks`: All blocks in the graph (order is preserved). * `node_shape_dtypes`: Maps every node class to its `jax.ShapeDtypeStruct` PyTree template. * `batch_shape`: Optional batch dimension(s) to prepend to every leaf. **Returns:** A list of PyTrees—one per _block_—whose leaves are `zeros(batch_shape + (len(block),) + leaf.shape)`. #### `thrml.verify_block_state(blocks: list[thrml.Block], states: list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']], node_shape_dtypes: Mapping[type[thrml.AbstractNode], PyTree[jax._src.core.ShapeDtypeStruct]], block_axis: int | None = None) -> None` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_management.py#L414-445) [¤](https://docs.thrml.ai/en/latest/api/block_management/#thrml.verify_block_state "Permanent link") Check that a state is what it should be given some blocks and node shape/dtypes. Passing incompatible state information into THRML functions can lead to unintended casting/other weird silent errors, so we should always check this. **Arguments:** * `blocks`: A list of Blocks. * `states`: A list of states to verify against blocks. * `node_shape_dtypes`: Maps every node class to its `jax.ShapeDtypeStruct` PyTree template. * `block_axis`: Index in the state batch shape at which to expect the block length. **Returns:** None. Raises RuntimeError if blocks and states are incompatible. --- # Interaction Groups - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/interaction/#interaction-groups) Interaction Groups[¤](https://docs.thrml.ai/en/latest/api/interaction/#interaction-groups "Permanent link") ============================================================================================================ #### `thrml.InteractionGroup` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/interaction.py#L9-66) [¤](https://docs.thrml.ai/en/latest/api/interaction/#thrml.InteractionGroup "Permanent link") Defines computational dependencies for conditional sampling updates. An `InteractionGroup` specifies information that is required to update the state of some subset of the nodes of a PGM during a block sampling routine. More concretely, when the state of the node at head\_nodes\[i\] is being updated, the sampler will receive the current state of the nodes at tail\_nodes\[k\]\[i\] for all k, and the ith element of each array in the Interaction PyTree (sliced along the first dimension). **Attributes:** * `head_nodes`: these are the nodes whose conditional updates should be affected by this InteractionGroup. * `tail_nodes`: these are the nodes whose state information is required to update `head_nodes`. * `interaction`: this specifies the parametric (independent of the state of the sampling program) required to update 'head\_nodes'. ##### `__init__(interaction: PyTree, head_nodes: thrml.Block, tail_nodes: list[thrml.Block])` [¤](https://docs.thrml.ai/en/latest/api/interaction/#thrml.InteractionGroup.__init__ "Permanent link") Create an `InteractionGroup`. An `InteractionGroup` implements a group of directed interactions between nodes in a PGM sampling program. **Arguments:** * `interaction`: A PyTree specifying the static information associated with the interaction. The first dimension of every Array in interaction must be equal to the length of `head_nodes`. * `head_nodes`: The nodes whose update is affected by the interaction. * `tail_nodes`: The groups of nodes whose state is required to update `head_nodes`. Each block in this list of blocks is intended to be parallel to `head_nodes`. i.e, to update the state of head\_nodes\[i\] during sampling we need state info about tail\_nodes\[k\]\[i\] for all values of k. --- # Factors - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/factor/#factors) Factors[¤](https://docs.thrml.ai/en/latest/api/factor/#factors "Permanent link") ================================================================================= #### `thrml.AbstractFactor` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/factor.py#L13-55) [¤](https://docs.thrml.ai/en/latest/api/factor/#thrml.AbstractFactor "Permanent link") A factor represents a batch of undirected interactions between sets of random variables. Concretely, this class implements a batch of factors defined over a bunch of parallel node groups. A single factor is defined over the nodes given by node\_groups\[k\]\[i\] for all k and a particular i. The defining trait of a factor is to produce InteractionGroups that affect each member of the factor in some way during the conditional updates of a block sampling program. As a user, you specify how this is done by implementing a concrete to\_interaction\_groups method for your child class. **Attributes:** * `node_groups`: the list of blocks that makes up this batch of factors. ##### `to_interaction_groups() -> list[thrml.InteractionGroup]` [¤](https://docs.thrml.ai/en/latest/api/factor/#thrml.AbstractFactor.to_interaction_groups "Permanent link") Compile a factor to a set of directed interactions. #### `thrml.WeightedFactor([thrml.AbstractFactor](https://docs.thrml.ai/en/latest/api/factor/#thrml.AbstractFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/factor.py#L58-77) [¤](https://docs.thrml.ai/en/latest/api/factor/#thrml.WeightedFactor "Permanent link") A factor that is parameterized by a weight tensor. The leading dimension of the weights tensor must be the same length as the batch dimension of the factor (i.e the number of nodes in each of the node\_groups). **Attributes:** * `weights`: the weight tensor. #### `thrml.FactorSamplingProgram([thrml.BlockSamplingProgram](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.BlockSamplingProgram) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/factor.py#L80-110) [¤](https://docs.thrml.ai/en/latest/api/factor/#thrml.FactorSamplingProgram "Permanent link") A sampling program built out of factors. This class simply breaks each factor passed to it down into interaction groups and uses them to build a BlockSamplingProgram. ##### `__init__(gibbs_spec: thrml.BlockGibbsSpec, samplers: list[thrml.AbstractConditionalSampler], factors: Sequence[thrml.AbstractFactor], other_interaction_groups: list[thrml.InteractionGroup])` [¤](https://docs.thrml.ai/en/latest/api/factor/#thrml.FactorSamplingProgram.__init__ "Permanent link") Create a FactorSamplingProgram. Thin wrapper over `BlockSamplingProgram`. **Arguments:** * `gibbs_spec`: A division of some PGM into free and clamped blocks. * `samplers`: The update rule to use for each free block in gibbs\_spec. * `factors`: The factors to use to build this sampling program. * `other_interaction_groups`: Other interaction groups to include in your program alongside what the factors produce. --- # Block Sampling - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/block_sampling/#block-sampling) Block Sampling[¤](https://docs.thrml.ai/en/latest/api/block_sampling/#block-sampling "Permanent link") ======================================================================================================= #### `thrml.BlockGibbsSpec([thrml.BlockSpec](https://docs.thrml.ai/en/latest/api/block_management/#thrml.BlockSpec) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L32-107) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.BlockGibbsSpec "Permanent link") A BlockGibbsSpec is a type of BlockSpec which contains additional information on free and clamped blocks. This entity also supports `SuperBlock`s, which are merely groups of blocks which are sampled at the same time algorithmically, but not programmatically. That is to say, superblock = (block1, block2) means that the states input to block1 and block2 are the same, but they are not executed at the same time. This may be because they are the same color on a graph, but require vastly different sampling methods such that JAX SIMD approaches are not feasible to parallelize them. A recurring theme in `thrml` is the importance of implicit indexing. One such example can be seen here. Because global states are created by concatenating lists of free and clamped blocks, providing the inputs in the same order as the blocks are defined is essential. This is almost always taken care of internally, but when writing custom functions or interfaces this is important to keep in mind. **Attributes:** * `free_blocks`: the list of free blocks (in order) * `sampling_order`: a list of `len(superblocks)` lists, where each `sampling_order[i]` is the index of `free_blocks` to sample. Sampling is done by iterating over this order and sampling each sublist of free blocks at the same algorithmic time. * `clamped_blocks`: the list of clamped blocks * `superblocks`: the list of superblocks ##### `__init__(free_super_blocks: Sequence[tuple[thrml.Block, ...] | thrml.Block], clamped_blocks: list[thrml.Block], node_shape_dtypes: Mapping[type[thrml.AbstractNode], PyTree[jax._src.core.ShapeDtypeStruct]] = {thrml.SpinNode: ShapeDtypeStruct(shape=(), dtype=bool), thrml.CategoricalNode: ShapeDtypeStruct(shape=(), dtype=uint8)})` [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.BlockGibbsSpec.__init__ "Permanent link") Create a Gibbs specification from free and clamped blocks. **Arguments:** * `free_super_blocks`: An ordered sequence where each element is either a single `Block`, or a tuple of blocks that must share the same global state when calling their individual samplers. * `clamped_blocks`: Blocks whose nodes stay fixed during sampling. * `node_shape_dtypes`: Mapping from node class to a PyTree of `jax.ShapeDtypeStruct`; identical to the argument in `BlockSpec`. #### `thrml.BlockSamplingProgram` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L116-263) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.BlockSamplingProgram "Permanent link") A PGM block-sampling program. This class encapsulates everything that is needed to run a PGM block sampling program in THRML. `per_block_interactions` and `per_block_interaction_active` are parallel to the free blocks in `gibbs_spec`, and their members are passed directly to a sampler when the state of the corresponding free block is being updated during a sampling program. `per_block_interaction_global_inds` and `per_block_interaction_global_slices` are also parallel to the free blocks, and are used to slice the global state of the program to produce the state information required to update the state of each block alongside the static information contained in the interactions. **Attributes:** * `gibbs_spec`: A division of some PGM into free and clamped blocks. * `samplers`: A sampler to use to update every free block in `gibbs_spec`. * `per_block_interactions`: All the interactions that touch each free block in `gibbs_spec`. * `per_block_interaction_active`: indicates which interactions are real and which interactions are not part of the model and have been added to pad data structures so that they can be rectangular. * `per_block_interaction_global_inds`: how to find the information required to update each block within the global state list * `per_block_interaction_global_slices`: how to slice each array in the global state list to find the information required to update each block ##### `__init__(gibbs_spec: thrml.BlockGibbsSpec, samplers: list[thrml.AbstractConditionalSampler], interaction_groups: list[thrml.InteractionGroup])` [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.BlockSamplingProgram.__init__ "Permanent link") Construct a `BlockSamplingProgram`. This code is the beating heart of THRML, and the chance that you should be modifying it or trying to understand it deeply are very low (as this would basically correspond to re-writing the library). This code takes in a set of information that implicitly defines a sampling program and manipulates it into a shape that is appropriate for practical vectorized block-sampling program. This involves reindexing, slicing, and often padding. **Arguments:** * `gibbs_spec`: A division of some PGM into free and clamped blocks. * `samplers`: The update rule to use for each free block in `gibbs_spec`. * `interaction_groups`: A list of `InteractionGroups` that define how the variables in your sampling program affect one another. #### `thrml.SamplingSchedule` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L407-424) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.SamplingSchedule "Permanent link") Represents a sampling schedule for a process. **Attributes:** * `n_warmup`: The number of warmup steps to run before collecting samples. * `n_samples`: The number of samples to collect. * `steps_per_sample`: The number of steps to run between each sample. #### `thrml.sample_blocks(key: Key[Array, ''], state_free: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], clamp_state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], program: thrml.BlockSamplingProgram, sampler_state: list[~_SamplerState]) -> tuple[list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], list[~_SamplerState]]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L342-381) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.sample_blocks "Permanent link") Perform one iteration of sampling, visiting every block. **Arguments:** * `key`: The JAX PRNG key. * `state_free`: The state of the free blocks. * `clamp_state`: The state of the clamped blocks. * `program`: The Gibbs program. * `sampler_state`: The state of the sampler. **Returns:** * Updated free-block state list and sampler-state list. #### `thrml.sample_single_block(key: Key[Array, ''], state_free: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], clamp_state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], program: thrml.BlockSamplingProgram, block: int, sampler_state: ~_SamplerState, global_state: list[PyTree] | None = None) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], '_State'], ~_SamplerState]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L269-339) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.sample_single_block "Permanent link") Samples a single block within a Gibbs sampling program based on the current states and program configurations. It extracts neighboring states, processes required data, and applies a sampling function to generate output samples. **Arguments:** * `key`: Pseudo-random number generator key to ensure reproducibility of sampling. * `state_free`: Current states of free blocks, representing the values to be updated during sampling. * `clamp_state`: Clamped states that remain fixed during the sampling process. * `program`: The Gibbs sampling program containing specifications, samplers, neighborhood information, and parameters. * `block`: Index of the block to be sampled in the current iteration. * `sampler_state`: The current state of the sampler that will be used to perform the update. * `global_state`: Optionally precomputed global state for the concatenated free and clamped blocks; when omitted the function constructs it internally. **Returns:** * Updated block state and sampler state for the specified block. #### `thrml.sample_with_observation(key: Key[Array, ''], program: thrml.BlockSamplingProgram, schedule: thrml.SamplingSchedule, init_chain_state: list[PyTree[Shaped[Array, 'nodes ?*state']]], state_clamp: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], observation_carry_init: ~ObserveCarry, f_observe: thrml.AbstractObserver) -> tuple[~ObserveCarry, list[PyTree[Shaped[Array, 'n_samples nodes ?*state']]]]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L427-506) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.sample_with_observation "Permanent link") Run the full chain and call an Observer after every recorded sample. **Arguments:** * `key`: RNG key. * `program`: The sampling program. * `schedule`: Warm-up length, number of samples, number of steps between samples. * `init_chain_state`: Initial free-block state. * `state_clamp`: Clamped-block state. * `observation_carry_init`: Initial carry handed to `f_observe`. * `f_observe`: Observer instance. **Returns:** * Tuple `(final_observer_carry, samples)` where `samples` is a PyTree whose leading axis has size `schedule.n_samples`. #### `thrml.sample_states(key: Key[Array, ''], program: thrml.BlockSamplingProgram, schedule: thrml.SamplingSchedule, init_state_free: list[PyTree[Shaped[Array, 'nodes ?*state']]], state_clamp: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], nodes_to_sample: list[thrml.Block]) -> list[PyTree[Shaped[Array, 'n_samples nodes ?*state']]]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/block_sampling.py#L509-536) [¤](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.sample_states "Permanent link") Convenience wrapper to collect state information for _nodes\_to\_sample_ only. Internally builds a [`thrml.StateObserver`](https://docs.thrml.ai/en/latest/api/observers/#thrml.StateObserver) , runs [`thrml.sample_with_observation`](https://docs.thrml.ai/en/latest/api/block_sampling/#thrml.sample_with_observation) , and returns a stacked tensor of shape `(schedule.n_samples, ...)`. --- # Ising Models - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/models/ising/#ising-models) Ising Models[¤](https://docs.thrml.ai/en/latest/api/models/ising/#ising-models "Permanent link") ================================================================================================= This module contains implementations of Ising models and spin systems. #### `thrml.models.IsingEBM([thrml.models.AbstractFactorizedEBM](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractFactorizedEBM) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L23-77) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingEBM "Permanent link") An EBM with the energy function, E(s)\=−β(∑i∈S1bisi+∑(i,j)∈S2Jijsisj) where S1 and S2 are the sets of biases and weights that make up the model, respectively. bi represents the bias associated with the spin si and Jij is a weight that couples si and sj. β is the usual temperature parameter. **Attributes:** * `nodes`: the nodes that have an associated bias (i.e S1) * `biases`: the bias associated with each node in `nodes`. * `edges`: the edges that have an associated weight (i.e S2) * `weights`: the weight associated with each pair of nodes in `edges`. * `beta`: the scalar temperature parameter for the model. ##### `__init__(nodes: list[thrml.AbstractNode], edges: list[tuple[thrml.AbstractNode, thrml.AbstractNode]], biases: Array, weights: Array, beta: Array)` [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingEBM.__init__ "Permanent link") Initialize an Ising EBM. **Arguments:** * `nodes`: List of nodes with associated biases * `edges`: List of edge pairs with associated weights * `biases`: Bias values for each node * `weights`: Weight values for each edge * `beta`: Temperature parameter #### `thrml.models.IsingSamplingProgram([thrml.FactorSamplingProgram](https://docs.thrml.ai/en/latest/api/factor/#thrml.FactorSamplingProgram) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L80-96) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingSamplingProgram "Permanent link") A very thin wrapper on FactorSamplingProgram that specializes it to the case of an Ising Model. ##### `__init__(ebm: thrml.models.IsingEBM, free_blocks: list[tuple[thrml.Block, ...] | thrml.Block], clamped_blocks: list[thrml.Block])` [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingSamplingProgram.__init__ "Permanent link") Initialize an Ising sampling program. **Arguments:** * `ebm`: The Ising EBM to sample from * `free_blocks`: List of super blocks that are free to vary * `clamped_blocks`: List of blocks that are held fixed #### `thrml.models.IsingTrainingSpec` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L99-128) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingTrainingSpec "Permanent link") Contains a complete specification of an Ising EBM that can be trained using sampling-based gradients. Defines sampling programs and schedules that allow for collection of the positive and negative phase samples required for Monte Carlo estimation of the gradient of the KL-divergence between the model and a data distribution. ##### `__init__(ebm: thrml.models.IsingEBM, data_blocks: list[thrml.Block], conditioning_blocks: list[thrml.Block], positive_sampling_blocks: list[tuple[thrml.Block, ...] | thrml.Block], negative_sampling_blocks: list[tuple[thrml.Block, ...] | thrml.Block], schedule_positive: thrml.SamplingSchedule, schedule_negative: thrml.SamplingSchedule)` [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.IsingTrainingSpec.__init__ "Permanent link") #### `thrml.models.hinton_init(key: Key[Array, ''], model: thrml.models.IsingEBM, blocks: list[thrml.Block[thrml.AbstractNode]], batch_shape: tuple[int]) -> list[Bool[Array, 'batch_size block_size']]` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L131-171) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.hinton_init "Permanent link") Initialize the blocks according to the marginal bias. Each binary unit i in a block is sampled independently as P(Si\=1)\=σ(βhi)\=11+e−βhi where hi is the bias of unit _i_ and β is the inverse-temperature scaling factor. See Hinton (2012) for a discussion of this initialization heuristic. Parameters: | Name | Type | Description | Default | | --- | --- | --- | --- | | `key` | `Key[Array, '']` | the JAX PRNG key to use | _required_ | | `model` | `thrml.models.IsingEBM` | the Ising model to initialize for | _required_ | | `blocks` | `list[thrml.Block[thrml.AbstractNode]]` | the blocks that are to be initialized | _required_ | | `batch_shape` | `tuple[int]` | the pre-pended dimension | _required_ | Returns: | Type | Description | | --- | --- | | `list[Bool[Array, 'batch_size block_size']]` | the initialized blocks | #### `thrml.models.estimate_moments(key: Key[Array, ''], first_moment_nodes: list[thrml.AbstractNode], second_moment_edges: list[tuple[thrml.AbstractNode, thrml.AbstractNode]], program: thrml.BlockSamplingProgram, schedule: thrml.SamplingSchedule, init_state: list[Array], clamped_data: list[Array])` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L174-215) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.estimate_moments "Permanent link") Estimates the first and second moments of an Ising model Boltzmann distribution via sampling. Parameters: | Name | Type | Description | Default | | --- | --- | --- | --- | | `key` | `Key[Array, '']` | the jax PRNG key | _required_ | | `first_moment_nodes` | `list[thrml.AbstractNode]` | the nodes that represent the variables we want to estimate the first moments of | _required_ | | `second_moment_edges` | `list[tuple[thrml.AbstractNode, thrml.AbstractNode]]` | the edges that connect the variables we want to estimate the second moments of | _required_ | | `program` | `thrml.BlockSamplingProgram` | the `BlockSamplingProgram` to be used for sampling | _required_ | | `schedule` | `thrml.SamplingSchedule` | the schedule to use for sampling | _required_ | | `init_state` | `list[Array]` | the variable values to use to initialize the sampling | _required_ | | `clamped_data` | `list[Array]` | the variable values to assign to the clamped nodes | _required_ | Returns: the first and second moment data #### `thrml.models.estimate_kl_grad(key: Key[Array, ''], training_spec: thrml.models.IsingTrainingSpec, bias_nodes: list[thrml.AbstractNode], weight_edges: list[tuple[thrml.AbstractNode, thrml.AbstractNode]], data: list[Array], conditioning_values: list[Array], init_state_positive: list[Array], init_state_negative: list[Array]) -> tuple` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ising.py#L218-296) [¤](https://docs.thrml.ai/en/latest/api/models/ising/#thrml.models.estimate_kl_grad "Permanent link") Estimate the KL-gradients of an Ising model with respect to its weights and biases. Uses the standard two-term Monte Carlo estimator of the gradient of the KL-divergence between an Ising model and a data distribution The gradients are: ΔW\=−β(⟨sisj⟩+−⟨sisj⟩−) Δb\=−β(⟨si⟩+−⟨si⟩−) Here, ⟨⋅⟩+ denotes an expectation under the _positive_ phase (data-clamped Boltzmann distribution) and ⟨⋅⟩− under the _negative_ phase (model distribution). Parameters: | Name | Type | Description | Default | | --- | --- | --- | --- | | `key` | `Key[Array, '']` | the JAX PRNG key | _required_ | | `training_spec` | `thrml.models.IsingTrainingSpec` | the Ising EBM for which to estimate the gradients | _required_ | | `bias_nodes` | `list[thrml.AbstractNode]` | the nodes for which to estimate the bias gradients | _required_ | | `weight_edges` | `list[tuple[thrml.AbstractNode, thrml.AbstractNode]]` | the edges for which to estimate the weight gradients | _required_ | | `data` | `list[Array]` | The data values to use for the positive phase of the gradient estimate. Each array has shape \[batch nodes\] | _required_ | | `conditioning_values` | `list[Array]` | values to assign to the nodes that the model is conditioned on. Each array has shape \[nodes\] | _required_ | | `init_state_positive` | `list[Array]` | initial state for the positive sampling chain. Each array has shape \[n\_chains\_pos batch nodes\] | _required_ | | `init_state_negative` | `list[Array]` | initial state for the negative sampling chain. Each array has shape \[n\_chains\_neg nodes\] | _required_ | Returns: the weight gradients and the bias gradients --- # Sampling Observers - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/observers/#sampling-observers) Sampling Observers[¤](https://docs.thrml.ai/en/latest/api/observers/#sampling-observers "Permanent link") ========================================================================================================== #### `thrml.AbstractObserver` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/observers.py#L21-63) [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.AbstractObserver "Permanent link") Interface for objects that inspect the sampling program while it is running. A concrete Observer is called once per block-sampling iteration and can maintain an arbitrary "carry" state across calls (e.g. running averages, histogram buffers, log-probs, etc.). ##### `init() -> PyTree` [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.AbstractObserver.init "Permanent link") Initialize the memory for the observer. Defaults to None. #### `thrml.StateObserver([thrml.AbstractObserver](https://docs.thrml.ai/en/latest/api/observers/#thrml.AbstractObserver) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/observers.py#L66-88) [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.StateObserver "Permanent link") Observer which logs the raw state of some set of nodes. **Attributes:** * `blocks_to_sample`: the list of `Block`s which the states are logged for ##### `__init__(blocks_to_sample: list[thrml.Block])` [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.StateObserver.__init__ "Permanent link") Initialize self. See help(type(self)) for accurate signature. #### `thrml.MomentAccumulatorObserver([thrml.AbstractObserver](https://docs.thrml.ai/en/latest/api/observers/#thrml.AbstractObserver) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/observers.py#L95-234) [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.MomentAccumulatorObserver "Permanent link") Observer that accumulates and updates the provided moments. It doesn't log any samples, and will only accumulate moments. Note that this observer does not scale the accumulated values by the number of times it was called. It simply records a running sum of a product of some state variables, ∑if(x1i)f(x2i)…f(xNi) **Attributes:** * `blocks_to_sample`: the blocks to accumulate the moments over. These are for constructing the final state, and aren't truly "blocks" in the algorithmic sense (they can be connected to each other). There is one block per node type. * `flat_nodes_list`: a list of all of the nodes in the moments (each occurring only once, so len(set(x)) = len(x)). * `flat_to_type_slices_list`: a list over node types in which each element is an array of indices of the `flat_node_list` which that type corresponds to * `flat_to_full_moment_slices`: a list over moment types in which each element is a 2D array, which matches the shape of the `moment_spec[i]` and of which each element is the index in the `flat_node_list`. * `f_transform`: the element-wise transformation f to apply to sample values before accumulation. ##### `__init__(moment_spec: Sequence[Sequence[Sequence[thrml.AbstractNode]]], f_transform: typing.Callable = )` [¤](https://docs.thrml.ai/en/latest/api/observers/#thrml.MomentAccumulatorObserver.__init__ "Permanent link") Create a MomentAccumulatorObserver. **Arguments:** * `moment_spec`: A 3 depth sequence. The first is a sequence over different moment types. A given moment type should have the same number of nodes in each moment. Then for each moment type, there is a sequence over moments. Each given moment is defined by a certain set of nodes. For example, to get the first and second moments on a simple o-o graph, it would be \[ \[(node1,), (node2,)\], \[(node1, node2)\] \] - `f_transform`: A function that takes in (state, blocks) and returns something with the same structure as state. This is used to apply functions to the samples before moments are computed. i.e this function defines a transformation of the state variable y\=f(x), such that the accumulated moments are of the form ⟨f(x1)f(x2)⟩. --- # Discrete Energy-Based Models - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#discrete-energy-based-models) Discrete Energy-Based Models[¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#discrete-energy-based-models "Permanent link") ======================================================================================================================================== This module contains implementations of discrete energy-based models. #### `thrml.models.DiscreteEBMFactor([thrml.models.EBMFactor](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.EBMFactor) , [thrml.WeightedFactor](https://docs.thrml.ai/en/latest/api/factor/#thrml.WeightedFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L60-194) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMFactor "Permanent link") Implements batches of energy function terms of the form s\_1 \* ... \* s\_M \* W\[c\_1, ..., c\_N\], where the s\_i are spin variables and the c\_i are categorical variables. No variable should show up twice in any given interaction. If this happens, the result of sampling from a model that includes the bad factor might not agree with the Boltzmann distribution. For example, the interaction w \* s\_1 \* s\_1 \* s\_2 would violate this rule because s\_1 shows up twice. To allow you to do something weird if you want to, this condition has not been enforced in the code. **Attributes:** * `spin_node_groups`: the node groups involved in the batch of factors that represent spin-valued random variables. * `categorical_node_groups`: the node groups involved in the batch of factors that represent categorical-valued random variables. * `weights`: the batch of weight tensors W associated with the factors we are implementing. `weights` should have leading dimension b, where b is number of nodes in each element of `spin_node_groups` and `categorical_node_groups`. This tensor has shape \[b, x\_1, ..., x\_N\] where b is the number of nodes in each block and N is the length of `categorical_node_groups`. * `is_spin`: a map that indicates if a given node type represents a spin-valued random variable or not. ##### `energy(global_state: list[Array], block_spec: thrml.BlockSpec)` [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMFactor.energy "Permanent link") Compute the energy associated with this factor. In this case, that is the sum of terms like s\_1 \* ... \* s\_M \* W\[c\_1, ..., c\_N\]. ##### `to_interaction_groups() -> list[thrml.InteractionGroup]` [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMFactor.to_interaction_groups "Permanent link") Produce interaction groups that implement this factor. In this case, we have to treat the spin and categorical node groups slightly differently. #### `thrml.models.DiscreteEBMInteraction` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L22-32) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMInteraction "Permanent link") An interaction that shows up when sampling from discrete-variable EBMs. **Attributes:** * `n_spin`: the number of spin states involved in the interaction. * `weights`: the weight tensor associated with this interaction. #### `thrml.models.SquareDiscreteEBMFactor([thrml.models.DiscreteEBMFactor](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L219-252) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SquareDiscreteEBMFactor "Permanent link") A discrete factor with a square interaction weight tensor (shape \[b, x, x, ..., x\]). If a discrete factor is square, the interaction groups corresponding to different choices of the head node blocks can be merged. This could yield smaller XLA programs and improved runtime performance via more efficient use of accelerators. ##### `to_interaction_groups() -> list[thrml.InteractionGroup]` [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SquareDiscreteEBMFactor.to_interaction_groups "Permanent link") Call the parent class to\_interaction\_groups, and merge the results. #### `thrml.models.SpinEBMFactor([thrml.models.SquareDiscreteEBMFactor](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SquareDiscreteEBMFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L255-259) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SpinEBMFactor "Permanent link") A `DiscreteEBMFactor` that involves only spin variables. #### `thrml.models.CategoricalEBMFactor([thrml.models.DiscreteEBMFactor](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.DiscreteEBMFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L262-266) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.CategoricalEBMFactor "Permanent link") A `DiscreteEBMFactor` that involves only categorical variables. #### `thrml.models.SquareCategoricalEBMFactor([thrml.models.SquareDiscreteEBMFactor](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SquareDiscreteEBMFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L269-273) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SquareCategoricalEBMFactor "Permanent link") A `DiscreteEBMFactor` that involves only categorical variables that also has a square weight tensor. #### `thrml.models.SpinGibbsConditional([thrml.BernoulliConditional](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.BernoulliConditional) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L309-348) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SpinGibbsConditional "Permanent link") A conditional update for spin-valued random variables that will perform a Gibbs sampling update given one or more `DiscreteEBMInteractions`. This function can be extended to handle a broader class of interactions via inheritance. Specifically, a child class can override the `compute_parameters` method defined here, compute contributions to γ from other types of interactions, and then call this method to take into account the contributions from `DiscreteEBMInteractions`. ##### `compute_parameters(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']]], sampler_state: None, output_sd: PyTree[jax._src.core.ShapeDtypeStruct]) -> PyTree` [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SpinGibbsConditional.compute_parameters "Permanent link") Compute the parameter γ of a spin-valued Bernoulli distribution given DiscreteEBMInteractions: γ\=∑is1i…sKiWi\[x1i,…,xMi\] where the sum over i is over all the `DiscreteEBMInteractions` seen by this function. #### `thrml.models.CategoricalGibbsConditional([thrml.SoftmaxConditional](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.SoftmaxConditional) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/discrete_ebm.py#L351-395) [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.CategoricalGibbsConditional "Permanent link") A conditional update for categorical random variables that will perform a Gibbs sampling update given one or more `DiscreteEBMInteractions`. This function can be extended to handle other interactions in the same way as [`thrml.models.SpinGibbsConditional`](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.SpinGibbsConditional) . **Attributes:** * `n_categories`: how many categories are involved in the softmax distribution this sampler will sample from. ##### `compute_parameters(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']]], sampler_state: None, output_sd: PyTree[jax._src.core.ShapeDtypeStruct]) -> PyTree` [¤](https://docs.thrml.ai/en/latest/api/models/discrete_ebm/#thrml.models.CategoricalGibbsConditional.compute_parameters "Permanent link") Compute the parameter θ of a softmax distribution given DiscreteEBMInteractions: θ\=∑is1i…sKiWi\[:,x1i,…,xMi\] where the sum over i is over all the `DiscreteEBMInteractions` seen by this function. --- # Energy-Based Models - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/models/ebm/#energy-based-models) Energy-Based Models[¤](https://docs.thrml.ai/en/latest/api/models/ebm/#energy-based-models "Permanent link") ============================================================================================================= This module contains implementations of energy-based models. #### `thrml.models.AbstractEBM` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ebm.py#L13-31) [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractEBM "Permanent link") Something that has a well-defined energy function (map from a state to a scalar). ##### `energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], blocks: list[thrml.Block]) -> Float[Array, '']` [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractEBM.energy "Permanent link") Evaluate the energy function of the EBM given some state information. **Arguments:** * `state`: The state for which to evaluate the energy function. Must be compatible with `blocks`. * `blocks`: Specifies how the information in `state` is organized. **Returns:** A scalar representing the energy value associated with `state`. #### `thrml.models.AbstractFactorizedEBM([thrml.models.AbstractEBM](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractEBM) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ebm.py#L50-84) [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractFactorizedEBM "Permanent link") An EBM that is made up of Factors, i.e., an EBM with an energy function like, E(x)\=∑iEi(x) where the sum over i is taken over factors. Child classes must define a property which returns a list of factors that substantiate the EBM. **Attributes:** * `node_shape_dtypes`: the shape/dtypes of the nodes involved in this EBM. Used to generate the BlockSpec that defines the global state that factors receive to compute energy. #### `thrml.models.FactorizedEBM([thrml.models.AbstractFactorizedEBM](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.AbstractFactorizedEBM) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ebm.py#L87-103) [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.FactorizedEBM "Permanent link") An EBM that is defined by a concrete list of factors. **Attributes:** * `_factors`: the list of factors that defines this EBM. ##### `__init__(factors: list[thrml.models.EBMFactor], node_shape_dtypes: Mapping[type[thrml.AbstractNode], PyTree[jax._src.core.ShapeDtypeStruct]] = {thrml.SpinNode: ShapeDtypeStruct(shape=(), dtype=bool), thrml.CategoricalNode: ShapeDtypeStruct(shape=(), dtype=uint8)})` [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.FactorizedEBM.__init__ "Permanent link") ##### `energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], blocks: list[thrml.Block]) -> Float[Array, '']` [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.FactorizedEBM.energy "Permanent link") #### `thrml.models.EBMFactor([thrml.AbstractFactor](https://docs.thrml.ai/en/latest/api/factor/#thrml.AbstractFactor) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/models/ebm.py#L34-47) [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.EBMFactor "Permanent link") A factor that defines an energy function. ##### `energy(global_state: list[Array], block_spec: thrml.BlockSpec) -> Float[Array, '']` [¤](https://docs.thrml.ai/en/latest/api/models/ebm/#thrml.models.EBMFactor.energy "Permanent link") Evaluate the energy function of the factor. **Arguments:** * `global_state`: The state information to use to evaluate the energy function. Is a global state of `block_spec`. * `block_spec`: The `BlockSpec` used to generate `global_state`. --- # Conditional Samplers - THRML [Skip to content](https://docs.thrml.ai/en/latest/api/conditional_samplers/#conditional-samplers) Conditional Samplers[¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#conditional-samplers "Permanent link") ========================================================================================================================= #### `thrml.AbstractConditionalSampler` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/conditional_samplers.py#L12-70) [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractConditionalSampler "Permanent link") Base class for all conditional samplers. A conditional sampler is used to update the state of a block of nodes during each iteration of a sampling algorithm. It takes in the states of all the neighbors and produces a sample for the current block of nodes. This can often be done exactly, but need not be. One could embed MCMC methods within this sampler (to do Metropolis within Gibbs, for example). ##### `sample(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]], sampler_state: ~_SamplerState, output_sd: PyTree[jax._src.core.ShapeDtypeStruct]) -> tuple[PyTree[Shaped[Array, 'nodes ?*state'], 'State'], ~_SamplerState]` [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractConditionalSampler.sample "Permanent link") Draw a sample from this conditional. If this sampler is involved in a block sampling program, this function is called every iteration to update the state of a block of nodes. **Arguments:** * `key`: A RNG key that the sampler can use to sample from distributions using `jax.random`. * `interactions`: A list of interactions that influence the result of this block update. Each interaction is a PyTree. Each array in the PyTree will have shape \[n, k, ...\], where n is the number of nodes in the block that is being updated and k is the maximum number of times any node in this block was detected as a head node for this interaction. * `active_flags`: A list of arrays of flags that is parallel to interactions. Each array indicates which instances of a given interaction are active for each node in the block. This array has shape \[n, k\], and is False if a given instance is inactive (which means that it should be ignored during the computation that happens in this function). * `states`: A list of PyTrees that is parallel to interactions, representing the sampling state information that is relevant to computing the influence of each interaction. Every array in each PyTree will have shape \[n, k, ...\]. * `sampler_state`: The current state of this sampler. Will be replaced by the second return from this function the next time it is called. * `output_sd`: A PyTree indicating the expected shape/dtype of the output of this function. **Returns:** A new state for the block of nodes, matching the template given by `output_sd`. #### `thrml.AbstractParametricConditionalSampler([thrml.AbstractConditionalSampler](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractConditionalSampler) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/conditional_samplers.py#L73-118) [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractParametricConditionalSampler "Permanent link") A conditional sampler that leverages a parameterized distribution. When `sample` is called, this sampler will first compute a set of parameters, and then use those parameters to draw a sample from some distribution. This workflow is frequently useful in practical cases; for example, to sample from a Gaussian, we can first compute a mean vector and covariance matrix using any procedure, and then draw a sample from the corresponding Gaussian distribution by appropriately transforming a vector of standard normal random variables. ##### `compute_parameters(key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[PyTree[Shaped[Array, 'nodes ?*state'], 'State']]], sampler_state: PyTree, output_sd: PyTree[jax._src.core.ShapeDtypeStruct]) -> PyTree` [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractParametricConditionalSampler.compute_parameters "Permanent link") Compute the parameters of the distribution. For a description of the arguments, see [`thrml.AbstractConditionalSampler.sample`](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractConditionalSampler.sample) #### `thrml.BernoulliConditional([thrml.AbstractParametricConditionalSampler](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractParametricConditionalSampler) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/conditional_samplers.py#L121-151) [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.BernoulliConditional "Permanent link") Sample from a bernoulli distribution. This sampler is designed to sample from a spin-valued bernoulli distribution: P(S\=s)∝eγs where S is a spin-valued random variable, s∈{−1,1}. The parameter γ must be computed by `compute_parameters`. #### `thrml.SoftmaxConditional([thrml.AbstractParametricConditionalSampler](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.AbstractParametricConditionalSampler) )` [](https://github.com/extropic-ai/thrml/blob/407691ec0d9767555ac57e48676f2ddd20c27cce/thrml/conditional_samplers.py#L154-184) [¤](https://docs.thrml.ai/en/latest/api/conditional_samplers/#thrml.SoftmaxConditional "Permanent link") Sample from a softmax distribution. This sampler samples from the standard softmax distribution: P(X\=k)∝eθk where X is a categorical random variable and θ is a vector that parameterizes the relative probabilities of each of the categories. --- # Spin Models in THRML - THRML [Skip to content](https://docs.thrml.ai/en/latest/examples/02_spin_models/#spin-models-in-thrml) Spin Models in THRML[¤](https://docs.thrml.ai/en/latest/examples/02_spin_models/#spin-models-in-thrml "Permanent link") ======================================================================================================================== Probabilistic computers that sample from graphical models defined over binary random variables are most natural to build using transistors, and therefore are of elevated interest to Extropic. As such, we've built some tooling into THRML that is dedicated to sampling from these binary PGMs and training machine learning models based on them. This notebook will walk through this functionality and show you how to use it. We specifically consider spin-valued EBMs with polynomial interactions. These models implement the probability distribution, P(x)∝e−E(x) E(x)\=−β(∑i∈S1Wi(1)si+∑(i,j)∈S2Wi,j(2)sisj+∑(i,j,k)∈S3Wi,j,k(3)sisjsk+…) Here, the si∈{−1,1} are spin variables that couple with each other via the W(k), which are scalars that represent the strengths of kth order interactions. Sk is the set of all interactions of order k. A model of this type that contains at most second-order interactions is called an Ising model or Boltzmann machine. Boltzmann machines are one of the original machine learning models, and their significance was recognized in 2024 with a Nobel prize in physics for John Hopfield and Geoffrey Hinton. Gibbs sampling defines a simple procedure for sampling from this type of model that is very hardware friendly. In particular, the Gibbs sampling update rule corresponding to the above energy function is, P(si\=1|snb(i))\=σ\[2γ\] γ\=Wi(1)+∑j∈S2\[i\]Wi,j(2)sj+∑(j,k)∈S3\[i\]Wi,j,k(3)sjsk+… where snb(i) are the spins that are neighbours of si, and Sk\[i\] is the members of Sk that contain i. From the above equation, we see that we can implement the Gibbs sampling update rule for a spin-valued model by computing simple functions of the neighbour states, multiply-accumulating the results, and then using them to generate an appropriately biased random bit. This can be done very efficiently using mixed signal (analog + digital) hardware; we flesh out a way to do this using only transistors on a modern process [in our recent paper](https://docs.thrml.ai/en/latest/examples/02_spin_models/denoising.paper) . Now that we understand the significance of this type of model, let's see how they can be sampled from using some of the tools built in to THRML. First, some imports, `import time import jax import dwave_networkx import jax.numpy as jnp import jax.random import matplotlib.pyplot as plt import networkx as nx import numpy as np` `from thrml.block_management import Block from thrml.block_sampling import sample_states, SamplingSchedule from thrml.models.discrete_ebm import SpinEBMFactor from thrml.models.ising import ( estimate_kl_grad, hinton_init, IsingEBM, IsingSamplingProgram, IsingTrainingSpec, ) from thrml.pgm import SpinNode` In this example, we will implement a quadratic binary model (Ising model). We will use DWave's "Pegasus" graph topology to allow us to directly compare the speed of our GPU-based sampler to results obtained using other hardware accelerators, `# make the graph using DWave's code graph = dwave_networkx.pegasus_graph(14) coord_to_node = {coord: SpinNode() for coord in graph.nodes} nx.relabel_nodes(graph, coord_to_node, copy=False)` `` Now we can define our model using the functionality exposed by `thrml.models.ising`. For the sake of this example, we will choose random values for the biases and weights W(1) and W(2), `nodes = list(graph.nodes) edges = list(graph.edges) seed = 4242 key = jax.random.key(seed) key, subkey = jax.random.split(key, 2) biases = jax.random.normal(subkey, (len(nodes),)) key, subkey = jax.random.split(key, 2) weights = jax.random.normal(subkey, (len(edges),)) beta = jnp.array(1.0) model = IsingEBM(nodes, edges, biases, weights, beta)` The `IsingEBM` class is simply a thin frontend that takes in your weights and biases and produces an appropriate set of `SpinEBMFactor`s, `[x.__class__ for x in model.factors]` `[thrml.models.discrete_ebm.SpinEBMFactor, thrml.models.discrete_ebm.SpinEBMFactor]` Now let's do some computation using our `IsingEBM`. Specifically, we are going to look at the tools THRML exposes for _training_ this type of model in the context of machine learning. In machine learning, the variables in an EBM are often segmented into "visible" variables (x) and "latent" variables (z). The visible variables represent the data, and the latent variables serve to increase the expressivity of the model. Given these latent variables, our EBMs model of the data is, P(x)∝∑ze−E(x,z) When training EBMs, one is often interested in minimizing the distributional distance between the EBM and some dataset. This can be done by iteratively updating the model parameters according to the gradient, ∇θD(Q(x)||P(x))\=EQ\[EP(z|x)\[∇θE\]−EP(z,x)\[∇θE\]\] Where D(Q||P) indicates the _KL-divergence_ between Q and P, which is a common measure of distributional distance in machine learning. Each of the two terms in this gradient can be estimated by sampling from the EBM. The first term is estimated by clamping the data nodes to a member of the dataset and sampling the latents. The second is estimated by sampling both the data and latent variables. We can leverage THRML for both of these computations. First, lets set up our block specifications for both the free and clamped sampling. First, lets choose some random subset of our nodes to represent the data, `n_data = 500 np.random.seed(seed) data_inds = np.random.choice(len(graph.nodes), n_data, replace=False) data_nodes = [nodes[x] for x in data_inds]` Now, lets compute the minimum coloring for the unclamped term in our gradient estimator, `coloring = nx.coloring.greedy_color(graph, strategy="DSATUR") n_colors = max(coloring.values()) + 1 free_coloring = [[] for _ in range(n_colors)] # form color groups for node in graph.nodes: free_coloring[coloring[node]].append(node) free_blocks = [Block(x) for x in free_coloring]` and the same for the clamped term, `# in this case we will just re-use the free coloring # you can always do this, but it might not be optimal # a graph without the data nodes graph_copy = graph.copy() graph_copy.remove_nodes_from(data_nodes) clamped_coloring = [[] for _ in range(n_colors)] for node in graph_copy.nodes: clamped_coloring[coloring[node]].append(node) clamped_blocks = [Block(x) for x in clamped_coloring]` We have now defined everything we need to calculate some gradients! We can set up a few more details and get to it, `# lets define some random "data" to use for our example # in real life this could be encoded images, text, video etc data_batch_size = 50 key, subkey = jax.random.split(key, 2) data = jax.random.bernoulli(subkey, 0.5, (data_batch_size, len(data_nodes))).astype(jnp.bool) # we will use the same sampling schedule for both cases schedule = SamplingSchedule(5, 100, 5) # convenient wrapper for everything you need for training training_spec = IsingTrainingSpec(model, [Block(data_nodes)], [], clamped_blocks, free_blocks, schedule, schedule) # how many parallel sampling chains to run for each term n_chains_free = data_batch_size n_chains_clamped = 1 # initial states for each sampling chain # THRML comes with simple code for implementing the hinton initialization, which is commonly used with boltzmann machines key, subkey = jax.random.split(key, 2) init_state_free = hinton_init(subkey, model, free_blocks, (n_chains_free,)) key, subkey = jax.random.split(key, 2) init_state_clamped = hinton_init(subkey, model, clamped_blocks, (n_chains_clamped, data_batch_size))` `# now for gradient estimation! # this function returns the gradient estimators for the weights and edges of our model, along with the moment data that was used to estimate them # the moment data is also returned in case you want to use it for something else in your training loop key, subkey = jax.random.split(key, 2) weight_grads, bias_grads, clamped_moments, free_moments = estimate_kl_grad( subkey, training_spec, nodes, # the nodes for which to compute bias gradients edges, # the edges for which to compute weight gradients [data], [], init_state_clamped, init_state_free, )` This function simply returns vectors for the weight and bias grads, `print(weight_grads) print(bias_grads)` `[ 0.7848 -0.33560008 -0.148 ... 0.00640005 -0.15759999 -0.01319999] [0.43279994 1.1767999 0.04360002 ... 0.01319999 0.18919998 0.14079998]` which can be used to train your model using whatever outer loop code you want! Because THRML is written in jax, it runs sampling programs very efficiently on GPUs and is competitive with the state of the art for sampling from sparse Ising models. Let's demonstrate that with a simple benchmark, Warning The following requires 8x GPUs. `from jax.sharding import PartitionSpec as P` `mesh = jax.make_mesh((8,), ("x",)) sharding = jax.sharding.NamedSharding(mesh, P("x")) timing_program = IsingSamplingProgram(model, free_blocks, []) timing_chain_len = 100 batch_sizes = [8, 80, 800, 8000, 64_000, 160_000, 320_000] times = [] flips = [] dofs = [] schedule = SamplingSchedule(timing_chain_len, 1, 1) call_f = jax.jit( jax.vmap(lambda k: sample_states(k, timing_program, schedule, [x[0] for x in init_state_free], [], [Block(nodes)])) ) for batch_size in batch_sizes: key, subkey = jax.random.split(key, 2) keys = jax.random.split(key, batch_size) keys = jax.device_put(keys, sharding) _ = jax.block_until_ready(call_f(keys)) start_time = time.time() _ = jax.block_until_ready(call_f(keys)) stop_time = time.time() times.append(stop_time - start_time) flips.append(timing_chain_len * len(nodes) * batch_size) dofs.append(batch_size * len(nodes))` `flips_per_ns = [x / (y * 1e9) for x, y in zip(flips, times)]` `fig, axs = plt.subplots() plt.title("Performance on 8xB200") axs.plot(dofs, flips_per_ns) axs.set_xscale("log") axs.set_xlabel("Parallel Degrees of Freedom") axs.set_ylabel("Flips/ns") plt.savefig("fps.png", dpi=300) plt.show()` ![](https://docs.thrml.ai/en/latest/examples/fps.png) You can compare your results to an FPGA implementation that bakes the sampling problem directly into hardware [here](https://arxiv.org/abs/2303.10728) (they get ~60 flips/ns). Note that despite our focus on quadratic models here, THRML comes with the ability to support spin interactions of arbitrary order out of the box. This ability can be accessed via `thrml.models.discrete_ebm.SpinEBMFactor`, `# this creates a cubic interaction s_1 * s_2 * s_3 between a subset of our nodes SpinEBMFactor([Block(nodes[:10]), Block(nodes[10:20]), Block(nodes[20:30])], jax.random.normal(key, (10,)))` `SpinEBMFactor( node_groups=[ Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())), Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())), Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())) ], weights=f32[10], spin_node_groups=[ Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())), Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())), Block(nodes=(SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode(), SpinNode())) ], categorical_node_groups=[], is_spin={thrml.pgm.SpinNode: True} )` That's about everything there is to know about binary EBMs in THRML! We hope you use these tools to help us gain a better understanding of how to most effectively use these powerful primitives in more advanced machine learning architectures. --- # All of THRML - THRML [Skip to content](https://docs.thrml.ai/en/latest/examples/01_all_of_thrml/#all-of-thrml) All of THRML[¤](https://docs.thrml.ai/en/latest/examples/01_all_of_thrml/#all-of-thrml "Permanent link") ========================================================================================================= THRML is a simple library for simulating probabilistic computers on GPUs. Concretely, THRML provides tools for GPU accelerating block sampling algorithms on sparse, heterogeneous probabilistic graphical models (PGMs) like the ones that Extropic hardware runs. The primary function of THRML is to be a scaffold that makes it much easier to implement any desired block sampling algorithm than it would be to do so from scratch. As such, this notebook will walk you through the main set of tools that THRML exposes that you can use in your own explorations. We will demonstrate the capabilities of THRML by using it to implement the Gibbs sampling algorithm for a Gaussian PGM. Gibbs sampling is obviously not a practical numerical method for Gaussian sampling in most cases, and should probably instead be handled using the [Cholesky decomposition](https://en.wikipedia.org/wiki/Cholesky_decomposition) . We implement it here solely for the purposes of demonstrating THRML, which in reality will be used to attack more complex problems that can't be treated analytically. Specifically, in the first part of this example we will consider a PGM that embodies the Gaussian distribution P(x)∝e−EG(x) Where the energy function EG(x) is, EG(x)\=12(x−μ)TA(x−μ) and A\=Σ−1, where Σ is the covariance matrix of the distribution. We can expand this and write the energy function as a sum of terms, EG(x)+C\=12∑iAiixi2+∑j\>iAijxixj+∑ibixi Where C is a constant independent of x, and b\=−μTA is a biasing vector. This form makes the graphical interpretation of the problem clear. Each of the variables xi can be represented by a node, and the nonzero matrix elements of A define edges between the nodes. We will use the Gibbs sampling algorithm to draw samples from our Gaussian distribution. To do this, we first identify the distribution of each xi conditioned on the rest of the graph, P(xi|xnb(i))∝e−Ei(xi,xnb(i)) Ei(xi,xnb(i))\=12Aiixi2+xi(∑j∈nb(i)Aijxj+bi) where nb(i) indicates the neighbours of node i, which in this case is all j such that Aij≠0. This form makes it clear that Gibbs sampling is local, i.e the state of each node is updated using only information about nodes that it is directly connected to. We can write this in a different form that makes it obvious that the conditional is Gaussian, Ei(xi,xnb(i))+D\=12(xi−mi)Aii(xi−mi) mi\=−(∑j∈nb(i)AijAiixj+biAii) where D is a constant independent of xi. The Gibbs sampling algorithm works by iteratively updating each of the xi according to this conditional distribution. In chromatic Gibbs sampling, nodes that belong to the same color group are updated in parallel. This "blocked" version is what we will implement here. With the math out of the way, we can proceed with the implementation of our sampling algorithm using THRML. First, let's get some imports out of the way: `import random from collections import defaultdict from typing import Hashable, Mapping import equinox as eqx import jax import jax.numpy as jnp import matplotlib.pyplot as plt import networkx as nx import numpy as np from jaxtyping import Array, Key, PyTree from sklearn.decomposition import PCA from sklearn.preprocessing import StandardScaler` `from thrml.block_management import Block from thrml.block_sampling import ( BlockGibbsSpec, BlockSamplingProgram, sample_states, sample_with_observation, SamplingSchedule, ) from thrml.conditional_samplers import ( _SamplerState, _State, AbstractConditionalSampler, ) from thrml.factor import AbstractFactor, FactorSamplingProgram from thrml.interaction import InteractionGroup from thrml.models.discrete_ebm import SpinEBMFactor, SpinGibbsConditional from thrml.observers import MomentAccumulatorObserver from thrml.pgm import AbstractNode` Next, we will define our graph. In THRML, nodes that represent random variables with different data types (binary, categorical, continuous, etc.) are identified using distinct classes that inherit from `AbstractNode`. For our problem we only have one type of node, which we will define now, `class ContinuousNode(AbstractNode): pass` We will now use the existing python graph library NetworkX to construct a grid graph of our nodes. `def generate_grid_graph( *side_lengths: int, ) -> tuple[ tuple[list[ContinuousNode], list[ContinuousNode]], tuple[list[ContinuousNode], list[ContinuousNode]], nx.Graph ]: G = nx.grid_graph(dim=side_lengths, periodic=False) coord_to_node = {coord: ContinuousNode() for coord in G.nodes} nx.relabel_nodes(G, coord_to_node, copy=False) for coord, node in coord_to_node.items(): G.nodes[node]["coords"] = coord # an aperiodic grid is always 2-colorable bicol = nx.bipartite.color(G) color0 = [n for n, c in bicol.items() if c == 0] color1 = [n for n, c in bicol.items() if c == 1] u, v = map(list, zip(*G.edges())) return (bicol, color0, color1), (u, v), G def plot_grid_graph( G: nx.Graph, bicol: Mapping[Hashable, int], ax: plt.Axes, *, node_size: int = 300, colors: tuple[str, str] = ("black", "orange"), **draw_kwargs, ): pos = {n: G.nodes[n]["coords"][:2] for n in G.nodes} node_colors = [colors[bicol[n]] for n in G.nodes] nx.draw( G, pos=pos, ax=ax, node_color=node_colors, node_size=node_size, edgecolors="k", linewidths=0.8, width=1.0, with_labels=False, **draw_kwargs, )` `colors, edges, g = generate_grid_graph(5, 5) all_nodes = colors[1] + colors[2] node_map = dict(zip(all_nodes, list(range(len(all_nodes))))) fig, axs = plt.subplots() plot_grid_graph(g, colors[0], axs)` ![img](https://docs.thrml.ai/en/latest/examples/_data/814f28ea7ae6490c88454ad80389adf4.png) The blue and orange nodes are the two color groups for our grid graph. Any nodes that are the same color will be sampled simultaneously during block sampling. With the graph in hand, we can fully define the distribution we want to sample from by choosing a corresponding inverse covariance matrix and mean vector, `# Fixed RNG seed for reproducibility seed = 4242 key = jax.random.key(seed) # diagonal elements of the inverse covariance matrix key, subkey = jax.random.split(key, 2) cov_inv_diag = jax.random.uniform(subkey, (len(all_nodes),), minval=1, maxval=2) # add an off-diagonal element to the inverse covariance matrix for each edge in the graph key, subkey = jax.random.split(key, 2) # make sure the covaraince matrix is PSD cov_inv_off_diag = jax.random.uniform(subkey, (len(edges[0]),), minval=-0.25, maxval=0.25) def construct_inv_cov(diag: Array, all_edges: tuple[list[ContinuousNode], list[ContinuousNode]], off_diag: Array): inv_cov = np.diag(diag) for n1, n2, cov in zip(*all_edges, off_diag): inv_cov[node_map[n1], node_map[n2]] = cov inv_cov[node_map[n2], node_map[n1]] = cov return inv_cov # construct a matrix representation of the inverse covariance matrix for convenience inv_cov_mat = construct_inv_cov(cov_inv_diag, edges, cov_inv_off_diag) inv_cov_mat_jax = jnp.array(inv_cov_mat) # mean vector key, subkey = jax.random.split(key, 2) mean_vec = jax.random.normal(subkey, (len(all_nodes),)) # bias vector b_vec = -1 * jnp.einsum("ij, i -> j", inv_cov_mat, mean_vec)` Now we can construct a program to sample from the distribution we just defined. All block sampling routines follow more or less the same set of steps 1. Divide your graph into two sets of blocks. The first set, the "free" blocks, will be updated during sampling. The second set, the "clamped" blocks, will have their nodes fixed to a constant value during sampling. This is often useful. For example, in the case of EBM sampling, this clamping allows for sampling from a distribution conditioned on the clamped nodes. 2. Iteratively update the states of your free blocks. This means: 1. Initialize the state of each of the free nodes 2. Update the state of each of the free nodes according to some rule. The update rule for each node is some function that takes in a set of parameters and the states of some subset of the other nodes in the graph, and returns an updated state for the node. 3. Make some observation of the current state of the program. This might mean simply writing down the state of some subset of the nodes, or it might mean computing some more complex observable. 4. Repeat steps 2 and 3 until a statisfactory number of observations have been made THRML lets you run any version of this procedure that you want while writing minimal amounts of new code. We have to define 3 main things to accomplish this: 1. A block specification: a division of our problem graph into free and clamped blocks 2. A set of interactions: these allow us to specify what information is required to compute the conditional updates for each node in our graph. 3. Conditional sampling rules: these specify how to update the state of each node in our graph given the interactions that are applicable to that node First, we will define a block spec for our problem. In our case, we simply want to sample each color group in sequence, and we won't be clamping any of the nodes, `# a Block is just a list of nodes that are all the same type # forcing the nodes in a Block to be of the same type is important for parallelization free_blocks = [Block(colors[1]), Block(colors[2])] # we won't be clamping anything here, but in principle this could be a list of Blocks just like above clamped_blocks = [] # every node in the program has to be assigned a shape and datatype (or PyTree thereof). # this is so THRML can build an internal "global" representation of the state of the sampling program using a small number of jax arrays node_shape_dtypes = {ContinuousNode: jax.ShapeDtypeStruct((), jnp.float32)} # our block specification spec = BlockGibbsSpec(free_blocks, clamped_blocks, node_shape_dtypes)` Now the interactions. Our PGM is of the undirected variety, which means that it can be described naturally using the language of [Factor Graphs](https://en.wikipedia.org/wiki/Factor_graph) . Deep knowledge of factor graphs and their nomenclature isn't necessary to use THRML; in this context, a Factor is simply an interaction between a set of variables that has no natural direction. Factor graphs can be viewed as hypergraphs where each factor represents a hyperedge connecting multiple variables. Hyperedges in factor graphs can connect any number of variables, allowing for natural representation of higher-order interactions. For example, a three-way interaction term like x1x2x3 in an energy function corresponds to a single hyperedge (factor) connecting three variable nodes. A nice thing about the Factor formalism is that in the context of Gibbs sampling, the conditional update rule for the ith node depends only on factors that involve xi. This means that given a set of factors for a graph, if we want to update the state of a given node, we only need to consider a small subset of all of the factors that are local to that node. In our case, our energy function can we written as a sum of a bunch of terms, each of which is associated with a factor. There are three distinct types of term in this sum, each of which is associated with a different type of factor: 1. Aiixi2 2. bixi 3. Aijxixj Each of these factors contributes to our conditional update rule in a different and consistent way. As such, in the context of algorithms like Gibbs sampling, Factors are defined by their ability to produce a set of directed interactions that effect the different nodes they involve in potentially different ways. In the case of our Gaussian sampling problem, our factors generate interactions that are either: 1. Linear: contribute terms to the energy function like cixi, where ci does not depend on the "head node" xi but may depend on some "tail nodes" xnb(i) 2. Quadratic: contribute terms to the energy function like dixi2, where in our case di is a constant independent of the state of the sampling program THRML implements these abstractions directly in code. The most primitive object is the `InteractionGroup`, which specifies what parametric and state information should be supplied to a given node to compute it's conditional update. An `InteractionGroup` is composed of a set of interaction parameters, a set of "head nodes", and sets of "tail nodes". The head nodes are the nodes whose conditional update is effected by the interaction, and the tail nodes specify which neighbouring node states are required to compute the conditional update. THRML also defines Factors via the `AbstractFactor` interface. In full generality, THRML defines a factor as anything that can be reduced to a set of `InteractionGroup`s. THRML also defines more specialized factors (like ones that define an energy), however we won't be using those here. We can use these objects to set up our sampling program, `# these are just arrays that we can identify by type, will be useful later class LinearInteraction(eqx.Module): """An interaction of the form $c_i x_i$.""" weights: Array class QuadraticInteraction(eqx.Module): """An interaction of the form $d_i x_i^2$.""" inverse_weights: Array # now we can set up our three different types of factors class QuadraticFactor(AbstractFactor): r"""A factor of the form $w \: x^2$""" # 1/A_{ii} inverse_weights: Array def __init__(self, inverse_weights: Array, block: Block): # in general, a factor is initialized via a list of blocks # these blocks should all have the same number of nodes, and represent groupings of nodes involved in the factor # for example, if a Factor involved 3 nodes, we would initialize it with 3 parallel blocks of equal length super().__init__([block]) # this array has shape [n], where n is the number of nodes in block self.inverse_weights = inverse_weights def to_interaction_groups(self) -> list[InteractionGroup]: # based on our conditional update rule, we can see that we need this to generate a Quadratic interaction with no tail nodes (i.e this interaction has no dependence on the neighbours of x_i) # we create an InteractionGroup that implements this interaction interaction = InteractionGroup( interaction=QuadraticInteraction(self.inverse_weights), head_nodes=self.node_groups[0], # no tail nodes in this case tail_nodes=[], ) return [interaction] class LinearFactor(AbstractFactor): r"""A factor of the form $w \: x$""" # b_i weights: Array def __init__(self, weights: Array, block: Block): super().__init__([block]) self.weights = weights def to_interaction_groups(self) -> list[InteractionGroup]: # follows the same pattern as previous, still no tail nodes return [ InteractionGroup(interaction=LinearInteraction(self.weights), head_nodes=self.node_groups[0], tail_nodes=[]) ] class CouplingFactor(AbstractFactor): # A_{ij} weights: Array def __init__(self, weights: Array, blocks: tuple[Block, Block]): # in this case our factor involves two nodes, so it is initialized with two blocks super().__init__(list(blocks)) self.weights = weights def to_interaction_groups(self) -> list[InteractionGroup]: # this factor produces interactions that impact both sets of nodes that it touches # i.e if this factor involves a term like w x_1 x_2, it should produce one interaction with weight w that has x_1 as a head node and x_2 as a tail node, # and another interaction with weight w that has x_2 as a head node and x_1 as a tail node # if we were sure that x_1 and x_2 were always the same type of node, the two interactions could be part of the same InteractionGroup # we won't worry about that here though return [ InteractionGroup(LinearInteraction(self.weights), self.node_groups[0], [self.node_groups[1]]), InteractionGroup(LinearInteraction(self.weights), self.node_groups[1], [self.node_groups[0]]), ]` Now the conditional update the rule. Here, we will define how the relevant interaction and state information should be used to produce an updated state in our iterative sampling algorithm. `class GaussianSampler(AbstractConditionalSampler): def sample( self, key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[_State]], sampler_state: _SamplerState, output_sd: PyTree[jax.ShapeDtypeStruct], ) -> tuple[Array, _SamplerState]: # this is where the rubber meets the road in THRML # this function gets called during block sampling, and must take in information about interactions and neighbour states and produce a state update # interactions, active_flags, and states are three parallel lists. # each item in interactions is a pytree, for which each array will have shape [n, k, ...]. # this is generated by THRML from the set of InteractionGroups that are used to create a sampling program # n is the number of nodes that we are updating in parallel during this call to sample # k is the maximum number of times any node in the block that is being updated shows up as a head node for this interaction # each item in active_flags is a boolean array with shape [n, k]. # this is padding that is generated internally by THRML based on the graphical structure of the model, # and serves to allow for heterogeneous graph sampling to be vectorized on accelerators that rely on homogeneous data structures # each item in states is a list of Pytrees that represents the state of the tail nodes that are relevant to this interaction. # for example, for an interaction with a single tail node that has a scalar dtype, states would be: # [[n, k],] bias = jnp.zeros(shape=output_sd.shape, dtype=output_sd.dtype) var = jnp.zeros(shape=output_sd.shape, dtype=output_sd.dtype) # loop through all of the available interactions and process them appropriately # here we are simply implementing the math of our conditional update rule for active, interaction, state in zip(active_flags, interactions, states): if isinstance(interaction, LinearInteraction): # if there are tail nodes, contribute w * x_1 * x_2 * ..., otherwise contribute w state_prod = jnp.array(1.0) if len(state) > 0: state_prod = jnp.prod(jnp.stack(state, -1), -1) bias -= jnp.sum(interaction.weights * active * state_prod, axis=-1) if isinstance(interaction, QuadraticInteraction): # this just sets the variance of the output distribution # there should never be any tail nodes var = active * interaction.inverse_weights var = var[..., 0] # there should only ever be one return (jnp.sqrt(var) * jax.random.normal(key, output_sd.shape)) + (bias * var), sampler_state def init(self) -> _SamplerState: return None` With all of the parts fully defined, we can now construct our sampling program `# our three types of factor lin_fac = LinearFactor(b_vec, Block(all_nodes)) quad_fac = QuadraticFactor(1 / cov_inv_diag, Block(all_nodes)) pair_quad_fac = CouplingFactor(cov_inv_off_diag, (Block(edges[0]), Block(edges[1]))) # an instance of our conditional sampler sampler = GaussianSampler() # the sampling program itself. Combines the three main components we just built prog = FactorSamplingProgram( gibbs_spec=spec, # one sampler for every free block in gibbs_spec samplers=[sampler, sampler], factors=[lin_fac, quad_fac, pair_quad_fac], other_interaction_groups=[], )` `FactorSamplingProgram` is a thin wrapper on the more generic `BlockSamplingProgram`. All `FactorSamplingProgram` does is convert all of the factors passed in into `InteractionGroups` and then use them to create a \`BlockSamplingProgram'. As such, prog is equivalent to prog\_2 in the following: `groups = [] for fac in [lin_fac, quad_fac, pair_quad_fac]: groups += fac.to_interaction_groups() prog_2 = BlockSamplingProgram(gibbs_spec=spec, samplers=[sampler, sampler], interaction_groups=groups)` Now we are finally ready to do some sampling! A sampling program in THRML simply repeatedly updates the state of each free block in the order they appear in the gibbs\_spec. After every iteration of the sampling algorithm, we may observe the state and write down some information that is relevant to the problem we are trying to solve. For example, if we wanted to extract samples from some subset of the nodes of our PGM, after each iteration we could simply memorize some subset of the current state. This functionality is provided by observers in THRML. For the purposes of this example, it would be prudent to check that our sampling program is working correctly. To do this, we will compute estimators of some first and second moments and verify that they match up with expected values from the theory. We will use the built-in `MomentAccumulatorObserver` to accomplish this. `# we will estimate the covariances for each pair of nodes connected by an edge and compare against theory # to do this we will need to estimate first moments and second moments second_moments = [(e1, e2) for e1, e2 in zip(*edges)] first_moments = [[(x,) for x in y] for y in edges] # this will accumulate products of the node state specified by first_moments and second_moments observer = MomentAccumulatorObserver(first_moments + [second_moments])` Now all that is left to do is specify a few more details about how the sampling should proceed. `# how many parallel sampling chains will we run? n_batches = 1000 schedule = SamplingSchedule( # how many iterations to do before drawing the first sample n_warmup=0, # how many samples to draw in total n_samples=10000, # how many steps to take between samples steps_per_sample=5, ) # construct the initial state of the iterative sampling algorithm init_state = [] for block in spec.free_blocks: key, subkey = jax.random.split(key, 2) init_state.append( 0.1 * jax.random.normal( subkey, ( n_batches, len(block.nodes), ), ) ) # RNG keys to use for each chain in the batch keys = jax.random.split(key, n_batches) # memory to hold our moment values init_mem = observer.init()` Now run the sampling: `# we use vmap to run a bunch of parallel sampling chains moments, _ = jax.vmap(lambda k, s: sample_with_observation(k, prog, schedule, s, [], init_mem, observer))( keys, init_state ) # Take a mean over the batch axis and divide by the total number of samples moments = jax.tree.map(lambda x: jnp.mean(x, axis=0) / schedule.n_samples, moments) # compute the covariance values from the moment data covariances = moments[-1] - (moments[0] * moments[1])` We can compare our covariance estimates to the real covariance matrix to see if we implemented our sampling routine correctly `cov = np.linalg.inv(inv_cov_mat) node_map = dict(zip(all_nodes, list(range(len(all_nodes))))) real_covs = [] for edge in zip(*edges): real_covs.append(cov[node_map[edge[0]], node_map[edge[1]]]) real_covs = np.array(real_covs) error = np.max(np.abs(real_covs - covariances)) / np.abs(np.max(real_covs)) print(error) assert error < 0.01` `0.0045360937` We achieve a really small error because we computed a ton of samples. If you reduce either the batch size or the number of samples collected by each chain this number will go up. That is everything you need to know to implement any type of PGM block sampling routine you want in THRML. However, you don't always have to do everything completely from scratch! THRML exposes a limited set of higher-level functionality fine-tuned to sampling problems that Extropic really cares about. Next, we will use some of these higher-level functions to implement a more complicated type of model that can't be sampled from using analytical techniques. In particular, we will implement sampling from a deep Gaussian-Bernoulli EBM. This type of model has the energy function E(x)\=EG(x)+EGB(x,s)+EB(s) where x is a vector of continuous values and s is a vector of _spins_, si∈{−1,1}. EG(x) is the Gaussian energy function defined in the previous section. EGB is an energy function that represents the interaction between the continuous and spin-valued variables, EGB(x,s)\=∑(i,j)∈SGBWijyixj where SGB is a set of edges connecting spin and continuous variables. EB is the spin energy function, EB(s)\=∑ibisi+∑j\>iJijsisj Just for fun, lets use a more complicated graph topology for this problem. We will stick with a grid, but we will add skip-connections that allow for non-nearest-neighbour interactions. We can once again use NetworkX to make this graph, `# first, define a new type of node class SpinNode(AbstractNode): pass` `# now, build a random grid out of spin and continuous nodes def make_random_typed_grid( rows: int, cols: int, seed: int, p_cont: float = 0.5, ): rng = random.Random(seed) # every time we make a node, flip a coin to decide what type it should be grid = [[ContinuousNode() if rng.random() < p_cont else SpinNode() for _ in range(cols)] for _ in range(rows)] # Parity-based 2-coloring bicol = {grid[r][c]: ((r + c) & 1) for r in range(rows) for c in range(cols)} # Separate by color and type colors_by_type = { 0: {SpinNode: [], ContinuousNode: []}, 1: {SpinNode: [], ContinuousNode: []}, } for r in range(rows): for c in range(cols): n = grid[r][c] color = bicol[n] colors_by_type[color][type(n)].append(n) return grid, colors_by_type grid, coloring = make_random_typed_grid(30, 30, seed)` `# now generate the edges to implement our desired skip-connected grid # we will use only odd-length edges (1, 3, 5, ...) so that our 2-coloring remains valid def build_skip_graph_from_grid( grid: list[list[AbstractNode]], skips: list[int], ): rows, cols = len(grid), len(grid[0]) # Build graph & annotate nodes with coords and type G = nx.Graph() for r in range(rows): for c in range(cols): n = grid[r][c] G.add_node(n, coords=(r, c)) # Edges sorted by edge length u_all = [] v_all = [] for k in skips: # vertical: (r, c) -> (r+k, c) for r in range(rows - k): r2 = r + k for c in range(cols): n1 = grid[r][c] n2 = grid[r2][c] u_all.append(n1) v_all.append(n2) G.add_edge(n1, n2) # horizontal: (r, c) -> (r, c+k) for r in range(rows): for c in range(cols - k): c2 = c + k n1 = grid[r][c] n2 = grid[r][c2] u_all.append(n1) v_all.append(n2) G.add_edge(n1, n2, skip=k) return (u_all, v_all), G edge_lengths = [1, 3, 5] edges, graph = build_skip_graph_from_grid(grid, edge_lengths)` Let's visualize this graph to understand what we just created. Since the graph is no longer planar, it will be cleanest to plot the local neighbourhood of particular nodes in our grid one at a time. `def plot_node_neighbourhood( grid, G: nx.Graph, center: Hashable, hops: int, ax: plt.Axes, ) -> None: rows, cols = len(grid), len(grid[0]) r, c = G.nodes[center]["coords"] # make a rectangular subgrid r0, r1 = max(0, r - hops), min(rows - 1, r + hops) c0, c1 = max(0, c - hops), min(cols - 1, c + hops) rect_nodes = {grid[i][j] for i in range(r0, r1 + 1) for j in range(c0, c1 + 1)} # collect the relevant edges by length edges_by_k = defaultdict(list) for v, ed in G[center].items(): k = int(ed.get("skip", 1)) edges_by_k[k].append((center, v)) # draw edges as arcs max_k = max(edges_by_k.keys(), default=1) curve_scale = 0.8 edge_width = 1.0 alpha = 1.0 def rad_for_edge(u, v, k): r1, c1 = G.nodes[u]["coords"] r2, c2 = G.nodes[v]["coords"] base = curve_scale * (k / max_k) # choose bend direction based on quadrant: if c1 == c2: sign = +1.0 if r2 < r1 else -1.0 # up vs down else: # horizontal edge sign = +1.0 if c2 > c1 else -1.0 # right vs left return sign * base # positions for plotting pos = {n: (G.nodes[n]["coords"][1], G.nodes[n]["coords"][0]) for n in rect_nodes | {center}} for i, k in enumerate(sorted(edges_by_k)): for u, v in edges_by_k[k]: nx.draw_networkx_edges( G, pos, edgelist=[(u, v)], ax=ax, edge_color="gray", width=edge_width, alpha=alpha, arrows=True, arrowstyle="-", connectionstyle=f"arc3,rad={rad_for_edge(u, v, k)}", ) # draw nodes cont_nodes = [n for n in rect_nodes if n.__class__ == ContinuousNode] spin_nodes = [n for n in rect_nodes if n.__class__ == SpinNode] node_size = 20.0 nx.draw_networkx_nodes(G, pos, nodelist=cont_nodes, node_color="black", node_shape="s", node_size=node_size, ax=ax) nx.draw_networkx_nodes(G, pos, nodelist=spin_nodes, node_color="orange", node_shape="o", node_size=node_size, ax=ax) # pick a few nodes in the grid to inspect centers = [grid[0][7], grid[10][10], grid[-1][-1]] fig, axs = plt.subplots(nrows=1, ncols=len(centers), figsize=(len(centers) * 5, 5)) for ax, center in zip(axs, centers): plot_node_neighbourhood(grid, graph, center, max(edge_lengths) + 1, ax)` ![img](https://docs.thrml.ai/en/latest/examples/_data/767df20a581b4aee83dab44814ffab57.png) This problem is clearly much more heterogeneous than what we were looking at before. Every node has a unique local neighbourhood, and is connected to a potentially different number of spin and continuous nodes. This makes working with this graph on an accelerator like a GPU tricky. As we will now see, THRML was specifically designed to handle this heterogeneity. With our graph in hand, let's set up our sampling program. We can re-use a lot of the work that we did in the simpler example. First, let's sort the nodes and edges by type. `# collect the different types of nodes spin_nodes = [] cont_nodes = [] for node in graph.nodes: if isinstance(node, SpinNode): spin_nodes.append(node) else: cont_nodes.append(node) # spin-spin interactions ss_edges = [[], []] # continuous-continuous interactions cc_edges = [[], []] # spin-continuous interactions sc_edges = [[], []] for edge in zip(*edges): if isinstance(edge[0], SpinNode) and isinstance(edge[1], SpinNode): ss_edges[0].append(edge[0]) ss_edges[1].append(edge[1]) elif isinstance(edge[0], ContinuousNode) and isinstance(edge[1], ContinuousNode): cc_edges[0].append(edge[0]) cc_edges[1].append(edge[1]) elif isinstance(edge[0], SpinNode): sc_edges[0].append(edge[0]) sc_edges[1].append(edge[1]) else: sc_edges[1].append(edge[0]) sc_edges[0].append(edge[1])` Now we can set up some interactions. For some of the factors, we will re-use our code from the first part of this example `# we will just randomize the weights key, subkey = jax.random.split(key, 2) cont_quad = QuadraticFactor(jax.random.uniform(subkey, (len(cont_nodes),), minval=2, maxval=3), Block(cont_nodes)) key, subkey = jax.random.split(key, 2) cont_linear = LinearFactor(jax.random.normal(subkey, (len(cont_nodes),)), Block(cont_nodes)) key, subkey = jax.random.split(key, 2) cont_coupling = CouplingFactor( jax.random.uniform(subkey, (len(cc_edges[0]),), minval=-1 / 10, maxval=1 / 10), (Block(cc_edges[0]), Block(cc_edges[1])), ) key, subkey = jax.random.split(key, 2) spin_con_coupling = CouplingFactor( jax.random.normal(subkey, (len(sc_edges[0]),)), (Block(sc_edges[0]), Block(sc_edges[1])) )` For the factors that involve only spin variables, we will use some built in functionality from THRML. THRML implements sampling functionality for arbitrary discrete-variable EBMs in `thrml.models.discrete_ebm` that we can apply to our problem. First, the spin factors, `key, subkey = jax.random.split(key, 2) spin_linear = SpinEBMFactor([Block(spin_nodes)], jax.random.normal(subkey, (len(spin_nodes),))) key, subkey = jax.random.split(key, 2) spin_coupling = SpinEBMFactor([Block(x) for x in ss_edges], jax.random.normal(subkey, (len(ss_edges[0]),)))` The Gaussian sampler we wrote for the first part will work our new problem as it is because it won't be seeing any new types of interactions. The Binary sampler built into THRML will have to be extended to handle our `LinearInteraction`. Luckily, it was designed with this kind of modification in mind. `class ExtendedSpinGibbsSampler(SpinGibbsConditional): def compute_parameters( self, key: Key, interactions: list[PyTree], active_flags: list[Array], states: list[list[_State]], sampler_state: _SamplerState, output_sd: PyTree[jax.ShapeDtypeStruct], ) -> PyTree: field = jnp.zeros(output_sd.shape, dtype=float) unprocessed_interactions = [] unprocessed_active = [] unprocessed_states = [] for interaction, active, state in zip(interactions, active_flags, states): # if its our new interaction, handle it if isinstance(interaction, LinearInteraction): state_prod = jnp.prod(jnp.stack(state, -1), -1) field -= jnp.sum(interaction.weights * active * state_prod, axis=-1) # if we haven't seen it, remember it else: unprocessed_interactions.append(interaction) unprocessed_active.append(active) unprocessed_states.append(state) # make the parent class deal with THRML-native interactions field -= super().compute_parameters( key, unprocessed_interactions, unprocessed_active, unprocessed_states, sampler_state, output_sd )[0] return field, sampler_state` This is all the work we need to do to sample from our new graph using THRML! All that is left is to set up our Block spec and run some sampling. `# tell THRML the shape and datatype of our new node new_sd = {SpinNode: jax.ShapeDtypeStruct(shape=(), dtype=jnp.bool)} # Our new graph is still two-colorable, however within each color there are two different types of node # this means that we can't make a single block to represent each color because all of the nodes within a block have to be of the same type # however, we might still want to ensure that the two blocks that represent each color group are sampled at the same "algorithmic" time # i.e even though we can't sample these blocks directly in parallel because they use different update rules, we want to make sure that they # receive the same state information # we can make this happen in THRML by passing in a list of tuples of blocks to BlockGibbsSpec instead of a list of Blocks # the blocks in each tuple will be sampled at the same algorithmic time blocks = [ (Block(coloring[0][SpinNode]), Block(coloring[0][ContinuousNode])), (Block(coloring[1][SpinNode]), Block(coloring[1][ContinuousNode])), ] block_spec = BlockGibbsSpec(blocks, [], node_shape_dtypes | new_sd) # now we can assemble our program # first, choose the right update rule for each block in the spec ber_sampler = ExtendedSpinGibbsSampler() samplers = [] for block in block_spec.free_blocks: if isinstance(block.nodes[0], SpinNode): samplers.append(ber_sampler) else: samplers.append(sampler) # collect all of our factors factors = [cont_quad, cont_linear, cont_coupling, spin_con_coupling, spin_linear, spin_coupling] program = FactorSamplingProgram(block_spec, samplers, factors, [])` Our program is doing a lot of work to pad out the interaction structure and make our sampling program GPU-compatible: `# let's look at an example of the padding program.per_block_interaction_active[0][0][0]` `Array([ True, True, True, True, True, False, False, False, False, False], dtype=bool)` Now we are ready to sample. In this case, we will simply observe the state of our nodes directly `batch_size = 50 schedule = SamplingSchedule( # how many iterations to do before drawing the first sample n_warmup=100, # how many samples to draw in total n_samples=300, # how many steps to take between samples steps_per_sample=15, ) # construct the initial state of the iterative sampling algorithm init_state = [] for block in block_spec.free_blocks: init_shape = ( batch_size, len(block.nodes), ) key, subkey = jax.random.split(key, 2) if isinstance(block.nodes[0], ContinuousNode): init_state.append(0.1 * jax.random.normal(subkey, init_shape)) else: init_state.append(jax.random.bernoulli(subkey, 0.5, init_shape)) key, subkey = jax.random.split(key, 2) keys = jax.random.split(subkey, batch_size) samples = jax.vmap(lambda k, i: sample_states(k, program, schedule, i, [], [Block(spin_nodes), Block(cont_nodes)]))( keys, init_state )` Let's visualize our samples. Our data is very high-dimensional, but we can use a PCA to try and get some idea of the structure of the distribution. `all_samples = jnp.concatenate(samples, axis=-1) pca = PCA(n_components=3) preproc_data = StandardScaler().fit_transform(jnp.reshape(all_samples, (-1, all_samples.shape[-1]))) transformed_data = pca.fit_transform(preproc_data)` `fig = plt.figure() ax = fig.add_subplot(111, projection="3d") sc = ax.scatter( transformed_data[:, 0], # PC1 transformed_data[:, 1], # PC2 transformed_data[:, 2], # PC3 s=50, alpha=0.8, ) ax.view_init(elev=-50, azim=280) plt.show()` ![img](https://docs.thrml.ai/en/latest/examples/_data/00275972b95e484cae9c408caa0b554a.png) This distribution is clearly non-Gaussian and complex, despite the random initialization. If you've made it to the end of this example and have been paying attention you are now ready to use THRML for your own research-grade problems! We are very excited to see what you make with it. --- # THRML 404 - Not found =============== --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... --- # Just a moment... docs.thrml.ai ============= Verifying you are human. This may take a few seconds. docs.thrml.ai needs to review the security of your connection before proceeding. Verification successful Waiting for docs.thrml.ai to respond... ---