mbi.MarginalLossFn

class mbi.MarginalLossFn(cliques: Sequence[tuple[str | int, ...]], loss_fn: Callable[[CliqueVector, Any], Array | ndarray | bool | number | bool | int | float | complex], data: Any = (), lipschitz: float = 1.0)

Bases: object

A loss function over the concatenated vector of marginals.

Separates the computation (loss_fn, static) from the captured data (data, dynamic JAX pytree leaves). This allows MarginalLossFn to be passed through jax.jit as a regular traced argument rather than a static one, so that different measurement values with the same structure reuse the same compiled program.

cliques

Cliques defining the scope of the marginals.

Type:

collections.abc.Sequence[tuple[str | int, …]]

loss_fn

A pure function (marginals, data) -> loss. This is treated as static metadata for JIT compilation.

Type:

collections.abc.Callable[[mbi.clique_vector.CliqueVector, Any], jax.Array | numpy.ndarray | numpy.bool | numpy.number | bool | int | float | complex]

data

Arbitrary pytree of arrays captured by loss_fn. This is traced through JIT and can change without recompilation.

Type:

Any

lipschitz

Lipschitz constant of the gradient. Defaults to 1.0.

Type:

float

Methods

__init__

Attributes

data
lipschitz
cliques
loss_fn

cliques: Sequence[tuple[str | int, ...]]

loss_fn: Callable[[CliqueVector, Any], Array | ndarray | bool | number | bool | int | float | complex]

data: Any = ()

lipschitz: float = 1.0

__call__(marginals: CliqueVector) → Array | ndarray | bool | number | bool | int | float | complex

Call self as a function.