mbi.marginal_oracles

Marginal oracles for computing marginals from graphical model potentials.

This module provides marginal oracles that compute marginals from the log-space potentials of a discrete graphical model. Three junction tree schedules are available:

• message_passing_implicit: Implicit-factor message passing using a configurable contraction function (default: einsum_fused).

• message_passing_hugin: Classic HUGIN algorithm with belief subtraction.

• message_passing_shafer_shenoy: Shafer-Shenoy algorithm with neighbor collection.

Use default_oracle() to automatically select the best oracle for your hardware and clique structure:

oracle = marginal_oracles.default_oracle(cliques=cliques, domain=domain)
marginals = oracle(potentials)

Functions

brute_force_marginals	Compute marginals from (log-space) potentials by materializing the full joint distribution.
bulk_variable_elimination	Compute the marginals of the graphical model with the given potentials.
calculate_many_marginals	Calculate marginals for all projections using belief propagation.
default_oracle	Select the best oracle for the given setting.
einsum_fused	Compute the log-space sum-product using custom_dot_general.
einsum_marginals	Compute marginals from (log-space) potentials by using einsum.
einsum_materialized	Compute the log-space sum-product by materializing the combined factor.
einsum_semistable	Compute the log-space sum-product via the exp-normalize trick.
kron_query	Compute a Kronecker-product query.
message_passing_hugin	HUGIN message passing with belief subtraction.
message_passing_implicit	Implicit-factor message passing using a contraction function.
message_passing_shafer_shenoy	Shafer-Shenoy message passing with neighbor collection.
message_passing_stable	HUGIN message passing with belief subtraction.
precompile_bulk_variable_elimination	Warm the JIT cache for bulk_variable_elimination asynchronously.
sum_product	Compute the sum-of-products of a list of Factors using einsum.
variable_elimination	Compute an out-of-model/unsupported marginal from the potentials.

Classes

MarginalOracle

Callable signature for stateless marginal oracle functions.

class mbi.marginal_oracles.MarginalOracle(*args, **kwargs)

Bases: Protocol

Callable signature for stateless marginal oracle functions.

A marginal oracle consumes log-space potentials of a graphical model and returns its marginals over the same set of cliques. Different oracles may have different time/space complexities and numerical stabilities.

Available oracles:

• message_passing_implicit: Implicit-factor message passing (configurable).

• message_passing_hugin: Classic HUGIN algorithm with belief subtraction.

• message_passing_shafer_shenoy: Shafer-Shenoy algorithm with neighbor collection.

• brute_force_marginals: Materializes the full joint distribution.

• einsum_marginals: Direct einsum (not recommended for large models).

mbi.marginal_oracles.sum_product(factors: list[~mbi.factor.Factor], dom: ~mbi.domain.Domain, einsum_fn: ~collections.abc.Callable = <function einsum>) → Factor

Compute the sum-of-products of a list of Factors using einsum.

Parameters:

• factors – A list of Factors.

• dom – The target domain of the output factor.

Returns:

sum_{S - D} prod_i F_i, where

• F_i = factors[i]

• D = dom

• S = union of domains of F_i

mbi.marginal_oracles.einsum_semistable(log_factors: list[~mbi.factor.Factor], dom: ~mbi.domain.Domain, einsum_fn: ~collections.abc.Callable = <function einsum>) → Factor

Compute the log-space sum-product via the exp-normalize trick.

Subtracts per-factor maxima for numerical stability, exponentiates, runs a standard multiply-sum einsum, then maps back to log space. Fast and memory-efficient (no super-factor materialization), but can produce NaN when factor magnitudes span a very wide range.

Parameters:

• log_factors – Factors in log space.

• dom – Target domain for the output.

• einsum_fn – Einsum implementation to use (default: jnp.einsum).

Returns:

A Factor over dom containing the result in log space.

mbi.marginal_oracles.einsum_materialized(log_factors: list[Factor], dom: Domain) → Factor

Compute the log-space sum-product by materializing the combined factor.

Combines all log-factors into a single super-factor (via addition in log-space), then reduces over the non-target dimensions with logsumexp. Numerically stable because no exp/log round-trip is needed, but may require O(product of all domain sizes) memory for the intermediate super-factor.

Parameters:

• log_factors – Factors in log space.

• dom – Target domain for the output.

Returns:

A Factor over dom in log space.

mbi.marginal_oracles.einsum_fused(log_factors: list[Factor], dom: Domain) → Factor

Compute the log-space sum-product using custom_dot_general.

Delegates to custom_einsum with add/logsumexp operations, allowing the XLA compiler to fuse the combine and reduce steps. This avoids both the exp/log round-trip of einsum_semistable and the super-factor materialization of einsum_materialized.

On GPU, XLA often fuses the operations resulting in both good stability and good performance.

Parameters:

• log_factors – Factors in log space.

• dom – Target domain for the output.

Returns:

A Factor over dom in log space.

mbi.marginal_oracles.message_passing_hugin(potentials: CliqueVector, total: Array | float = 1.0, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), return_messages: bool = False) → CliqueVector | tuple[CliqueVector, dict]

HUGIN message passing with belief subtraction.

Expands potentials to super-cliques and uses belief subtraction for message computation. Unstable when potentials contain -inf.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• jtree – Pre-computed junction tree (optional).

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

• return_messages – If True, return (marginals, messages) where messages is a dict mapping (clique_i, clique_j) to the log-space message Factor sent from i to j.

mbi.marginal_oracles.message_passing_shafer_shenoy(potentials: CliqueVector, total: Array | float = 1.0, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), return_messages: bool = False) → CliqueVector | tuple[CliqueVector, dict]

Shafer-Shenoy message passing with neighbor collection.

Expands potentials to super-cliques and collects messages from all neighbors except the target. More stable than HUGIN for -inf potentials.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• jtree – Pre-computed junction tree (optional).

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

• return_messages – If True, return (marginals, messages) where messages is a dict mapping (clique_i, clique_j) to the log-space message Factor sent from i to j.

mbi.marginal_oracles.message_passing_implicit(potentials: CliqueVector, total: Array | float = 1.0, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), contraction: Callable = <function einsum_materialized>, return_messages: bool = False) → CliqueVector | tuple[CliqueVector, dict]

Implicit-factor message passing using a contraction function.

Keeps potentials in their original factored form and computes messages via the given contraction function. Most memory-efficient — never materializes super-clique tables.

The default contraction (einsum_materialized) is the fastest on GPU. For better -inf tolerance, use einsum_fused.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• jtree – Pre-computed junction tree (optional).

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

• contraction – Contraction function for log-space sum-product.

• return_messages – If True, return (marginals, messages) where messages is a dict mapping (clique_i, clique_j) to the log-space message Factor sent from i to j.

mbi.marginal_oracles.message_passing_fast(potentials: CliqueVector, total: Array | float = 1.0, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), contraction: Callable = <function einsum_semistable>, return_messages: bool = False) → CliqueVector | tuple[CliqueVector, dict]

Implicit-factor message passing using a contraction function.

Keeps potentials in their original factored form and computes messages via the given contraction function. Most memory-efficient — never materializes super-clique tables.

The default contraction (einsum_materialized) is the fastest on GPU. For better -inf tolerance, use einsum_fused.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• jtree – Pre-computed junction tree (optional).

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

• contraction – Contraction function for log-space sum-product.

• return_messages – If True, return (marginals, messages) where messages is a dict mapping (clique_i, clique_j) to the log-space message Factor sent from i to j.

mbi.marginal_oracles.message_passing_stable(potentials: CliqueVector, total: Array | float = 1.0, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), return_messages: bool = False) → CliqueVector | tuple[CliqueVector, dict]

HUGIN message passing with belief subtraction.

Expands potentials to super-cliques and uses belief subtraction for message computation. Unstable when potentials contain -inf.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• jtree – Pre-computed junction tree (optional).

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

• return_messages – If True, return (marginals, messages) where messages is a dict mapping (clique_i, clique_j) to the log-space message Factor sent from i to j.

mbi.marginal_oracles.default_oracle(cliques: Sequence[tuple[str | int, ...]] | None = None, domain: Domain | None = None, backend: str | None = None, has_constraints: bool = True) → MarginalOracle

Select the best oracle for the given setting.

Chooses based on hardware backend, clique structure, and whether the potentials may contain -inf entries:

• CPU: message_passing_shafer_shenoy (has_constraints=True) or message_passing_hugin (has_constraints=False). Implicit is not used on CPU because the XLA compiler is less effective there.

• GPU/TPU, memory-limited: message_passing_implicit when the total beliefs memory would exceed ~40% of estimated device HBM. Implicit keeps factors in their original low-dimensional form and only materializes one super-clique at a time, using orders of magnitude less memory than Hugin/SS which materialize all super-clique arrays simultaneously.

• GPU/TPU, huge but few super-cliques: message_passing_implicit when individual super-cliques exceed 50M entries but there are fewer than 30 nodes. With few messages, implicit moves less data per message by contracting small input factors rather than operating on huge belief arrays.

• GPU/TPU, has_constraints=True (default): message_passing_shafer_shenoy (numerically robust).

• GPU/TPU, has_constraints=False: message_passing_hugin (faster when -inf is absent).

The heuristic is based on systematic A100 benchmarks across varied graph structures (cardinality 4-256, bandwidth 1-3, 32-256 attributes). When super-cliques are large but the junction tree has many nodes, Hugin and Shafer-Shenoy outperform implicit by 30-40% because their uniform add/logsumexp kernel pattern fuses better in XLA over many sequential messages. Implicit only wins when individual super-cliques are very large (>50M entries) and the tree is small (<30 nodes), or when Hugin/SS would exceed device memory.

Parameters:

• cliques – Cliques of the graphical model. Used to estimate max clique size when domain is also provided.

• domain – Domain of the graphical model. Used with cliques to estimate max clique size.

• backend – JAX backend string (‘cpu’, ‘gpu’, ‘tpu’). If None, uses jax.default_backend().

• has_constraints – Whether potentials may contain -inf entries (e.g. from deterministic constraints or structural zeros). When True (default), Shafer-Shenoy is preferred for numerical robustness. When False, Hugin may be used for better performance.

Returns:

A callable satisfying the MarginalOracle protocol.

Example:

oracle = default_oracle(cliques=model.cliques, domain=model.domain)
marginals = oracle(potentials)

mbi.marginal_oracles.brute_force_marginals(potentials: CliqueVector, total: Array | float = 1, *, constraints: Sequence[Constraint] = ()) → CliqueVector

Compute marginals from (log-space) potentials by materializing the full joint distribution.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

mbi.marginal_oracles.einsum_marginals(potentials: CliqueVector, total: Array | float = 1, einsum_fn: Callable = <function einsum>, *, constraints: Sequence[Constraint] = ()) → CliqueVector

Compute marginals from (log-space) potentials by using einsum.

This is a “brute-force” approach and is not recommended in practice.

Parameters:

• potentials – Potentials of a graphical model.

• total – Normalization constant.

• einsum_fn – Einsum function to use.

• constraints – Structural constraints. Raises ValueError if non-empty.

mbi.marginal_oracles.variable_elimination(potentials: CliqueVector, clique: tuple[str, ...], total: float = 1.0, evidence: dict[str | int, int | Array] | None = None, *, constraints: Sequence[Constraint] = ()) → Factor

Compute an out-of-model/unsupported marginal from the potentials.

Parameters:

• potentials – The (log-space) potentials of a Graphical Model.

• clique – The subset of attributes whose marginal you want.

• total – The normalization factor.

• evidence – Mapping from attribute names to observed scalar int values.

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

Returns:

The marginal defined over the domain of the input clique, where each entry is non-negative and sums to the input total.

mbi.marginal_oracles.precompile_bulk_variable_elimination(domain: Domain, potential_cliques: Sequence[tuple[str | int, ...]], marginal_queries: Sequence[tuple[str | int, ...]], *, constraints: Sequence[Constraint] = ()) → Future

Warm the JIT cache for bulk_variable_elimination asynchronously.

Fire and forget. Only requires domain and clique structure.

Parameters:

• domain – The Domain of the graphical model.

• potential_cliques – The cliques that potentials are defined over.

• marginal_queries – A list of cliques to obtain marginals for.

• constraints – Structural constraints.

mbi.marginal_oracles.bulk_variable_elimination(potentials: CliqueVector, marginal_queries: Sequence[tuple[str | int, ...]], total: float = 1.0, *, constraints: Sequence[Constraint] = ()) → CliqueVector

Compute the marginals of the graphical model with the given potentials.

Unlike other marginal oracles, which only compute marginals for cliques in the potentials vector, this function can compute arbitrary marginals from an arbitrary model. Both runtime and compilation time can be expensive when there are a large number of marginal queries. This function compiles and runs variable_elimination for one query at a time, using parallelism and asyncronous computation do do the compilation in the background, while running variable_eliminatoin sequentially one query at a time.

Parameters:

• potentials – The (log-space) potentials of a Graphical Model.

• marginal_queries – A list of cliques to obtain marginals for.

• total – The normalization factor.

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

Returns:

A CliqueVector with the marginals computed over the specified cliques.

mbi.marginal_oracles.calculate_many_marginals(potentials: ~mbi.clique_vector.CliqueVector, marginal_queries: ~collections.abc.Sequence[tuple[str | int, ...]], total: float = 1.0, belief_propagation_oracle: ~mbi.marginal_oracles.MarginalOracle = <PjitFunction of <function message_passing_hugin>>, *, constraints: ~collections.abc.Sequence[~mbi.constraint.Constraint] = ()) → CliqueVector

Calculate marginals for all projections using belief propagation.

Implements Algorithm from section 10.3 in Koller and Friedman. This method may be faster than calling variable_elimination many times. Note: this implementation is experimental, and further work may be needed to optimize it. Contributions are welcome.

Parameters:

• potentials – Potentials of a graphical model.

• marginal_queries – a list of cliques whose marginals are desired.

• constraints – Structural constraints folded into potentials as -inf/0 factors before inference.

Returns:

A CliqueVector, where each defined over the list of input marginal_queries.

mbi.marginal_oracles.kron_query(potentials: CliqueVector, query_factors: dict[str, Array], total: float = 1, suffix: str = '_answer', *, constraints: Sequence[Constraint] = ()) → Factor

Compute a Kronecker-product query.

Parameters:

• potentials – Potentials of a graphical model.

• query_factors – Mapping from attribute names to query matrices.

• total – Normalization constant.

• suffix – Suffix for the answer attributes.

• constraints – Structural constraints. Raises ValueError if non-empty.