Skip to content

Block Management#

Blocks are disjoint subsets of nodes. BlockSpec handles the index bookkeeping — mapping between block-local and global state representations, padding variable-size blocks for JAX array stacking, and verifying consistency.

hamon.Block #

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 hamon, this means that the nodes must all be of the same type.

Attributes:

  • nodes: the tuple of nodes that this block contains

hamon.BlockSpec #

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

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[hamon.block_management.Block], node_shape_dtypes: typing.Mapping[typing.Type[hamon.pgm.AbstractNode], PyTree[jax.ShapeDtypeStruct]]) -> None #

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 Blocks that this specification operates on
  • node_shape_dtypes: the mapping of node types to their structures. This should be a pytree of jax.ShapeDtypeStructs.

hamon.block_state_to_global(block_state: list[PyTree[Shaped[Array, 'nodes ?*state'], State]], spec: BlockSpec) -> list[PyTree[Shaped[Array, 'nodes_global ?*state'], _GlobalState]] #

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 hamon.BlockSpec that defines the mapping.

Returns:

A list whose length equals len(spec.global_sd_order)—the stacked global state.

hamon.get_node_locations(nodes: Block, spec: BlockSpec) -> tuple[int, Int[Array, nodes]] #

Locate a contiguous set of nodes inside the global state.

Arguments:

Returns:

Tuple (sd_index, positions) where

  • sd_index is the position inside the global list returned by hamon.block_state_to_global, and
  • positions is a 1D array with the indices each node occupies inside that particular PyTree.

hamon.from_global_state(global_state: list[PyTree[Shaped[Array, 'nodes_global ?*state'], _GlobalState]], spec_from: BlockSpec, blocks_to_extract: list[hamon.block_management.Block]) -> list[PyTree[Shaped[Array, 'nodes ?*state'], State]] #

Extract the states for a subset of blocks from a global state.

Arguments:

Returns:

A list with one element per blocks_to_extract—each element is a PyTree with exactly len(block) nodes in its leading dimension.

hamon.make_empty_block_state(blocks: list[hamon.block_management.Block], node_shape_dtypes: typing.Mapping[typing.Type[hamon.pgm.AbstractNode], PyTree[jax.ShapeDtypeStruct]], batch_shape: typing.Optional[tuple] = None) -> list[PyTree[Shaped[Array, 'nodes ?*state'], State]] #

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

hamon.verify_block_state(blocks: list[hamon.block_management.Block], states: list[PyTree[Shaped[Array, 'nodes ?*state'], State]], node_shape_dtypes: typing.Mapping[typing.Type[hamon.pgm.AbstractNode], PyTree[jax.ShapeDtypeStruct]], block_axis: typing.Optional[int] = None) -> None #

Check that a state is what it should be given some blocks and node shape/dtypes.

Passing incompatible state information into hamon 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.