mbi.extensions package

Submodules

mbi.extensions.message_passing module

Extension message-passing oracles with efficient constraint handling.

Alternative implementations of the Shafer-Shenoy and implicit message-passing oracles from mbi.marginal_oracles. Key differences:

More efficient constraint handling. For deterministic (mapping) constraints, messages are routed through the constraint in O(|fine|) time via coarsen/refine operations, avoiding the full O(|fine| × |coarse|) potential materialization used by the core oracles.
AI-written, less well-tested. These implementations were developed with AI assistance and have narrower test coverage than the core oracles. They are kept separate until they mature enough to replace the originals.

Usage:

from mbi.extensions.message_passing import shafer_shenoy, implicit

mbi.extensions.message_passing.coarsen(factor, constraint)[source]: Log-space fine-to-coarse aggregation via segment-logsumexp.

mbi.extensions.message_passing.default_oracle(cliques: tuple[tuple[str, ...], ...] | None = None, domain: Domain | None = None, backend: str | None = None) → MarginalOracle[source]

Select the best constraint-aware oracle for the given setting.

When constraints are present, this oracle uses the extensions implementations which route messages through deterministic constraints efficiently. When no constraints are passed at call time, behavior is identical to the core marginal_oracles.default_oracle.

The selection heuristic mirrors the core oracle:

CPU: shafer_shenoy (XLA compiler less effective on CPU).
GPU/TPU, large cliques (>= 1M): implicit.
GPU/TPU, small cliques: shafer_shenoy.

mbi.extensions.message_passing.implicit(potentials: CliqueVector, total: float = 1, jtree: Graph | None = None, *, constraints: Sequence[Constraint] = (), contraction: Callable = <function einsum_materialized>) → CliqueVector[source]

Hybrid message passing: implicit for standard cliques, SS for constraints.

Uses the memory-efficient implicit contraction for cliques that do not involve deterministic constraints, and falls back to Shafer-Shenoy style constraint routing for pure and embedded constraint cliques. General (valid/invalid) constraints are folded in as materialized potentials.

Parameters:

potentials – The (log-space) potentials of a graphical model.
total – The normalization factor.
jtree – An optional junction tree.
constraints – Structural constraints to handle.
contraction – Contraction function for log-space sum-product.

Returns:

The marginals of the graphical model.

mbi.extensions.message_passing.project_to_coarse(factor, constraint)[source]: Probability-space fine-to-coarse aggregation via segment-sum.

mbi.extensions.message_passing.refine(factor, constraint)[source]: Log-space coarse-to-fine scatter via indexing.

mbi.extensions.message_passing.shafer_shenoy(potentials: CliqueVector, total: float = 1, jtree: Graph | None = None, constraints: Sequence[Constraint] = ()) → CliqueVector[source]

Compute marginals using Shafer-Shenoy with constraint-aware shortcuts.

Deterministic (mapping) constraints are routed through efficient O(|fine|) coarsen/refine operations. General (valid/invalid) constraints are folded into potentials as materialized -inf/0 factors.

Parameters:

potentials – The (log-space) potentials of a graphical model.
total – The normalization factor.
jtree – An optional junction tree that defines the message passing order.
constraints – Structural constraints to handle.

Returns:

The marginals of the graphical model.

mbi.extensions.mixture_of_products module

Mixture-of-products estimation for discrete distributions.

Inspired by RAP (https://arxiv.org/abs/2103.06641) and RAP^{softmax} (https://arxiv.org/abs/2106.07153). Instead of parameterizing via a graphical model with potentials, the distribution is represented as a mixture of K product distributions optimized via gradient descent with optax.

This is a cleaned-up, optimized replacement for mbi.experimental.mixture_of_products.

class mbi.extensions.mixture_of_products.MixtureOfProducts(_logits: dict[str | int, Array], domain: Domain, total: float)[source]

Bases: object

A discrete distribution as a mixture of product distributions.

domain: Domain

property num_components: int: Number of mixture components K.

property products: dict[str | int, Array]: Per-attribute categorical probabilities (softmax of logits).

project(attrs: str | int | Sequence[str | int]) → Factor[source]: Compute the marginal over attrs via einsum.

classmethod random(domain: Domain, total: float, num_components: int = 100, seed: int = 0, scale: float = 0.25) → MixtureOfProducts[source]: Initialize with random logits.

supports(attrs: str | int | Sequence[str | int]) → bool[source]: Any subset of domain attributes is supported.

synthetic_data(rows=None) → Dataset[source]: Generate synthetic data via randomized rounding.

total: float

class mbi.extensions.mixture_of_products.MixtureOfProductsEstimator(*, marginal_oracle: marginal_oracles.MarginalOracle | None = None, optimizer: optax.GradientTransformation, num_components: int = 100, seed: int = 0)[source]

Bases: Estimator

Estimates a distribution as a mixture of product distributions.

Inspired by RAP (https://arxiv.org/abs/2103.06641) and RAP^{softmax} (https://arxiv.org/abs/2106.07153). Instead of parameterizing via a graphical model with potentials, the distribution is represented as a mixture of K product distributions optimized via gradient descent.

num_components

Number of mixture components K.

Type:: int

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

seed

Random seed for parameter initialization.

Type:: int

num_components: int = 100

optimizer: GradientTransformation

seed: int = 0

class mbi.extensions.mixture_of_products.MixtureOfProductsState(model: MixtureOfProducts, opt_state: optax.OptState)[source]

Bases: NamedTuple

Optimization state for MixtureOfProductsEstimator.

Create new instance of MixtureOfProductsState(model, opt_state)

model: MixtureOfProducts: Alias for field number 0

opt_state: optax.OptState: Alias for field number 1

mbi.extensions.precompute_marginals module

Precompute all marginals from a dataset using JIT-compiled bincount.

Computes a CliqueVector of marginals from raw data by JIT-compiling a jnp.bincount kernel per unique shape. Three optimizations reduce compilation overhead for heterogeneous domains:

Power-of-2 bucketing: Domain sizes are rounded up to the next power of 2, so attributes with similar cardinalities share a compiled program.
Canonical shape ordering: Within each clique, attributes are sorted so that padded sizes are non-decreasing. This halves the number of unique shapes for 2-way marginals.
Background precompilation: Unique shapes are sorted by workload (descending number of cliques), and all compilations are submitted upfront so they overlap with execution.

mbi.extensions.precompute_marginals.precompute_marginals(dataset: Dataset | JaxDataset, cliques: Sequence[tuple[str | int, ...]], *, transpose: bool = True) → CliqueVector[source]

Compute all marginals from a dataset using JIT-compiled bincount.

Compiles one jnp.bincount program per unique padded shape (after power-of-2 bucketing and canonical reordering) and executes it for every clique that maps to that shape. All compilations are submitted to a background thread upfront so they overlap with execution.

Accepts both Dataset (numpy-backed) and JaxDataset (JAX-backed). Numpy arrays are converted to JAX arrays internally.

Parameters:

dataset – The dataset to compute marginals from.
cliques – Cliques (tuples of attribute names) to compute.
transpose – If True (default), transpose each marginal so its axes match the original clique attribute order. Set to False to skip the transpose and keep the canonical sorted order.

Returns:

A CliqueVector mapping each clique to its marginal Factor.

mbi.extensions.reweighted_dataset module

Reweighted dataset estimation for discrete distributions.

Inspired by PMW^Pub (https://arxiv.org/abs/2102.08598). Given a seed set of records, the distribution is represented by learning a weight for each record to minimize the marginal loss. The weights are parameterized in log-space with softmax normalization to guarantee non-negativity.

class mbi.extensions.reweighted_dataset.ReweightedDatasetEstimator(seed_data: Dataset, optimizer: optax.GradientTransformation, *, marginal_oracle: marginal_oracles.MarginalOracle | None = None)[source]

Bases: Estimator

Estimates a distribution by reweighting seed records.

Inspired by PMW^Pub (https://arxiv.org/abs/2102.08598). Given a seed set of records, the distribution is represented by learning a weight for each record to minimize the marginal loss.

seed_data

A Dataset of seed records to reweight.

Type:: mbi.dataset.Dataset

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

optimizer: GradientTransformation

seed_data: Dataset

class mbi.extensions.reweighted_dataset.ReweightedDatasetState(log_weights: jax.Array, opt_state: optax.OptState)[source]

Bases: NamedTuple

Optimization state for ReweightedDatasetEstimator.

Create new instance of ReweightedDatasetState(log_weights, opt_state)

log_weights: jax.Array: Alias for field number 0

opt_state: optax.OptState: Alias for field number 1

mbi.extensions.synthetic_data module

JAX-accelerated synthetic data generation with async precompilation.

mbi.extensions.synthetic_data.precompile(domain: Domain, cliques: Sequence[tuple[str | int, ...]], rows: int, constraints: Sequence[Constraint] = ()) → Future[source]

Warm the JIT cache for synthetic_data asynchronously.

Only requires domain and clique structure — both available after measurement selection. Fire and forget; synthetic_data benefits from whatever has compiled so far.

No concrete arrays are allocated (except a PRNG key); all inputs use abstract ShapeDtypeStruct values via Factor.abstract so the compiler sees shapes and dtypes without materializing memory.

Parameters:

domain – The Domain over which the model is defined.
cliques – The cliques of the model (known after measurement selection).
rows – Number of records that will be generated.
constraints – Structural constraints the model is fit under. Their cliques are added to the plan so the warmed JIT cache matches the constraint-aware synthetic_data call.

mbi.extensions.synthetic_data.synthetic_data(model: MarkovRandomField, rows: int, seed: int = 0) → Dataset[source]

Generate synthetic data from a fitted model using Gumbel rounding.

Output marginals match the model within ±2 counts per cell. If precompile was called earlier with matching domain, cliques, and rows, the JIT cache will be warm and this call is fast.

Parameters:

model – A fitted MarkovRandomField.
rows – Number of records to generate.
seed – Random seed for the JAX PRNG.

Returns:

A Dataset whose marginals closely match the model.

Module contents

Extensions for mbi providing alternative estimation approaches.

class mbi.extensions.MixtureOfProducts(_logits: dict[str | int, Array], domain: Domain, total: float)[source]

Bases: object

A discrete distribution as a mixture of product distributions.

domain: Domain

property num_components: int: Number of mixture components K.

property products: dict[str | int, Array]: Per-attribute categorical probabilities (softmax of logits).

project(attrs: str | int | Sequence[str | int]) → Factor[source]: Compute the marginal over attrs via einsum.

classmethod random(domain: Domain, total: float, num_components: int = 100, seed: int = 0, scale: float = 0.25) → MixtureOfProducts[source]: Initialize with random logits.

supports(attrs: str | int | Sequence[str | int]) → bool[source]: Any subset of domain attributes is supported.

synthetic_data(rows=None) → Dataset[source]: Generate synthetic data via randomized rounding.

total: float

class mbi.extensions.MixtureOfProductsEstimator(*, marginal_oracle: marginal_oracles.MarginalOracle | None = None, optimizer: optax.GradientTransformation, num_components: int = 100, seed: int = 0)[source]

Bases: Estimator

Estimates a distribution as a mixture of product distributions.

Inspired by RAP (https://arxiv.org/abs/2103.06641) and RAP^{softmax} (https://arxiv.org/abs/2106.07153). Instead of parameterizing via a graphical model with potentials, the distribution is represented as a mixture of K product distributions optimized via gradient descent.

num_components

Number of mixture components K.

Type:: int

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

seed

Random seed for parameter initialization.

Type:: int

num_components: int = 100

optimizer: GradientTransformation

seed: int = 0

class mbi.extensions.ReweightedDatasetEstimator(seed_data: Dataset, optimizer: optax.GradientTransformation, *, marginal_oracle: marginal_oracles.MarginalOracle | None = None)[source]

Bases: Estimator

Estimates a distribution by reweighting seed records.

Inspired by PMW^Pub (https://arxiv.org/abs/2102.08598). Given a seed set of records, the distribution is represented by learning a weight for each record to minimize the marginal loss.

seed_data

A Dataset of seed records to reweight.

Type:: mbi.dataset.Dataset

optimizer

An optax optimizer (e.g. optax.adam(0.1)).

Type:: optax._src.base.GradientTransformation

optimizer: GradientTransformation

seed_data: Dataset

mbi.extensions.precompile(domain: Domain, cliques: Sequence[tuple[str | int, ...]], rows: int, constraints: Sequence[Constraint] = ()) → Future[source]

Warm the JIT cache for synthetic_data asynchronously.

Only requires domain and clique structure — both available after measurement selection. Fire and forget; synthetic_data benefits from whatever has compiled so far.

No concrete arrays are allocated (except a PRNG key); all inputs use abstract ShapeDtypeStruct values via Factor.abstract so the compiler sees shapes and dtypes without materializing memory.

Parameters:

domain – The Domain over which the model is defined.
cliques – The cliques of the model (known after measurement selection).
rows – Number of records that will be generated.
constraints – Structural constraints the model is fit under. Their cliques are added to the plan so the warmed JIT cache matches the constraint-aware synthetic_data call.

mbi.extensions.precompute_marginals(dataset: Dataset | JaxDataset, cliques: Sequence[tuple[str | int, ...]], *, transpose: bool = True) → CliqueVector[source]

Compute all marginals from a dataset using JIT-compiled bincount.

Compiles one jnp.bincount program per unique padded shape (after power-of-2 bucketing and canonical reordering) and executes it for every clique that maps to that shape. All compilations are submitted to a background thread upfront so they overlap with execution.

Accepts both Dataset (numpy-backed) and JaxDataset (JAX-backed). Numpy arrays are converted to JAX arrays internally.

Parameters:

dataset – The dataset to compute marginals from.
cliques – Cliques (tuples of attribute names) to compute.
transpose – If True (default), transpose each marginal so its axes match the original clique attribute order. Set to False to skip the transpose and keep the canonical sorted order.

Returns:

A CliqueVector mapping each clique to its marginal Factor.

mbi.extensions.synthetic_data(model: MarkovRandomField, rows: int, seed: int = 0) → Dataset[source]

Generate synthetic data from a fitted model using Gumbel rounding.

Output marginals match the model within ±2 counts per cell. If precompile was called earlier with matching domain, cliques, and rows, the JIT cache will be warm and this call is fast.

Parameters:

model – A fitted MarkovRandomField.
rows – Number of records to generate.
seed – Random seed for the JAX PRNG.

Returns:

A Dataset whose marginals closely match the model.