THRML

Energy-based models

Energy-based models define a distribution through an energy function. THRML factorizes that energy so block Gibbs can sample it.

AbstractEBMclass
AbstractEBM()

Something that has a well-defined energy function (map from a state to a scalar).

energymethod
energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], blocks: list[Block]) -> Float[Array, '']

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.

AbstractFactorizedEBMclass
AbstractFactorizedEBM(node_shape_dtypes: Mapping[Type[AbstractNode], PyTree[jax.ShapeDtypeStruct]] = {SpinNode: ShapeDtypeStruct(shape=(), dtype=bool), CategoricalNode: ShapeDtypeStruct(shape=(), dtype=uint8)})

An EBM that is made up of Factors, i.e., an EBM with an energy function like,

$$\mathcal{E}(x) = \sum_i \mathcal{E}^i(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.
node_shape_dtypesattribute
node_shape_dtypes: Mapping[Type[AbstractNode], PyTree[jax.ShapeDtypeStruct]]
energymethod
energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], blocks: list[Block]) -> Float[Array, '']

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.

factorsproperty
factors

A concrete implementation of this class must define this method that returns a list of factors that substantiate the EBM.

FactorizedEBMclass
FactorizedEBM(factors: list[EBMFactor], node_shape_dtypes: Mapping[Type[AbstractNode], PyTree[jax.ShapeDtypeStruct]] = {SpinNode: ShapeDtypeStruct(shape=(), dtype=bool), CategoricalNode: ShapeDtypeStruct(shape=(), dtype=uint8)})

An EBM that is defined by a concrete list of factors.

Attributes:

  • _factors: the list of factors that defines this EBM.
node_shape_dtypesattribute
node_shape_dtypes: Mapping[Type[AbstractNode], PyTree[jax.ShapeDtypeStruct]]
energymethod
energy(state: list[PyTree[Shaped[Array, 'nodes ?*state'], '_State']], blocks: list[Block]) -> Float[Array, '']

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.

factorsproperty
factors
EBMFactorclass
EBMFactor(node_groups: list[Block])

A factor that defines an energy function.

node_groupsattribute
node_groups: list[Block]
energymethod
energy(global_state: list[Array], block_spec: BlockSpec) -> Float[Array, '']

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.
to_interaction_groupsmethod
to_interaction_groups() -> list[InteractionGroup]

Compile a factor to a set of directed interactions.