"""Factor representation for discrete probability distributions."""
from __future__ import annotations
import dataclasses
from collections.abc import Callable, Sequence
from typing import Literal
from jax.typing import ArrayLike
import jax
import jax.numpy as jnp
import numpy as np
from .domain import Attribute, Domain
def _safe_logsumexp(a: jax.Array, axis: tuple[int, ...]) -> jax.Array:
# Double-where keeps the backward softmax finite (0 instead of exp(-inf - -inf))
# when a constraint forces an entire slice along `axis` to -inf.
all_neg_inf = jnp.isneginf(jnp.max(a, axis=axis, keepdims=True))
safe_a = jnp.where(all_neg_inf, 0.0, a)
lse = jax.scipy.special.logsumexp(safe_a, axis=axis)
return jnp.where(jnp.squeeze(all_neg_inf, axis=axis), -jnp.inf, lse)
[docs]
@jax.tree_util.register_dataclass
@dataclasses.dataclass(frozen=True) # pytype: disable=invalid-function-definition
class Factor:
"""Represents a factor defined over a discrete domain.
A factor can be thought of as a potential function or an unnormalized
probability distribution over a set of discrete variables defined by a
`Domain` object. It maps each configuration of the domain to a value.
Attributes:
domain (Domain): The discrete domain over which the factor is defined.
values (jax.Array): A JAX array containing the factor's values. The shape
of this array matches the shape specified by the `domain`.
Supported Operations:
- Creation: `zeros`, `ones`, `random` for creating factors.
- Reshaping: `transpose`, `expand` for modifying the domain/shape.
- Aggregation: `sum`, `logsumexp`, `project` for marginalizing attributes.
- Element-wise: `exp`, `log`, `normalize` for value transformations.
- Binary Ops: `+`, `-`, `*`, `/`, `dot` for combining factors.
Example Usage:
>>> from mbi import Domain # Needed for doctest context
>>> domain = Domain.fromdict({'X': 2, 'Y': 3})
>>> factor = Factor.ones(domain)
>>> print(factor.domain)
Domain(X: 2, Y: 3)
"""
domain: Domain = jax.tree.static()
values: jax.Array
# Constructors
[docs]
@classmethod
def zeros(cls, domain: Domain) -> Factor:
"""Creates a Factor object with all values initialized to zero."""
return cls(domain, jnp.zeros(domain.shape))
[docs]
@classmethod
def ones(cls, domain: Domain) -> Factor:
"""Creates a Factor object with all values initialized to one."""
return cls(domain, jnp.ones(domain.shape))
[docs]
@classmethod
def random(cls, domain: Domain) -> Factor:
"""Creates a Factor object with random values (uniform 0-1)."""
return cls(domain, jnp.asarray(np.random.rand(*domain.shape)))
[docs]
@classmethod
def abstract(
cls,
domain: Domain,
sharding: jax.sharding.Sharding | None = None,
) -> Factor:
"""Creates a shape/dtype-only Factor for ahead-of-time compilation."""
# Match zeros/ones' default float dtype so abstract-traced programs share
# the JIT cache with real runs (a fixed dtype would force recompiles).
dtype = jnp.result_type(float)
return cls(
domain,
# ShapeDtypeStruct stands in for a concrete Array under eval_shape.
jax.ShapeDtypeStruct(domain.shape, dtype, sharding=sharding), # pyrefly: ignore[bad-argument-type]
)
# Reshaping operations
[docs]
def transpose(self, attrs: Sequence[Attribute]) -> Factor:
"""Rearranges the factor's axes according to the new attribute order."""
if set(attrs) != set(self.domain.attributes):
raise ValueError("attrs must be same as domain attributes")
newdom = self.domain.project(attrs)
ax = newdom.axes(self.domain.attributes)
values = jnp.moveaxis(self.values, range(len(ax)), ax)
return Factor(newdom, values)
[docs]
def expand(self, domain: Domain) -> Factor:
"""Expands the factor's domain to include new attributes."""
if not domain.contains(self.domain):
raise ValueError("Expanded domain must contain domain.")
dims = len(domain) - len(self.domain)
values = self.values.reshape(self.domain.shape + tuple([1] * dims))
ax = domain.axes(self.domain.attributes)
values = jnp.moveaxis(values, range(len(ax)), ax)
values = jnp.broadcast_to(values, domain.shape)
return Factor(domain, values)
# Functions that aggregate along some subset of axes
[docs]
def aggregate(
self, fn: Callable, attrs: Sequence[Attribute] | None = None
) -> Factor:
"""Aggregates values along specified attribute axes using ``fn``."""
if fn is jax.scipy.special.logsumexp:
fn = _safe_logsumexp
attrs = self.domain.attributes if attrs is None else attrs
axes = self.domain.axes(attrs)
if not axes:
return self
values = fn(self.values, axis=axes)
newdom = self.domain.marginalize(attrs)
return Factor(newdom, values)
[docs]
def max(self, attrs: Sequence[Attribute] | None = None) -> Factor:
"""Computes the maximum value along specified attribute axes."""
return self.aggregate(jnp.max, attrs)
[docs]
def sum(self, attrs: Sequence[Attribute] | None = None) -> Factor:
"""Computes the sum along specified attribute axes."""
return self.aggregate(jnp.sum, attrs)
[docs]
def logsumexp(self, attrs: Sequence[Attribute] | None = None) -> Factor:
"""Computes the log-sum-exp along specified attribute axes."""
return self.aggregate(_safe_logsumexp, attrs)
[docs]
def project(
self,
attrs: Attribute | Sequence[Attribute],
log: bool = False,
reduction_fn: Callable | None = None,
) -> "Factor":
"""Computes the marginal distribution by aggregating out other attributes."""
if isinstance(attrs, (str, int)):
attrs = (attrs,)
if reduction_fn is None:
reduction_fn = _safe_logsumexp if log else jnp.sum
marginalized = self.domain.marginalize(attrs).attributes
return self.aggregate(reduction_fn, marginalized).transpose(attrs)
[docs]
def slice(self, evidence: dict[Attribute, int | jax.Array]) -> "Factor":
"""Slice the factor by fixing attributes to scalar values.
Args:
evidence: Mapping from attribute names to observed integer values.
Returns:
A new Factor with the evidence attributes removed.
"""
slices: list[slice | int | jax.Array] = [slice(None)] * len(self.domain)
relevant = [e for e in evidence if e in self.domain.attributes]
for attr in relevant:
val = evidence[attr]
if hasattr(val, "ndim") and val.ndim > 0:
raise ValueError(
"Array-valued evidence is not supported (got"
f" {type(val).__name__} for '{attr}'). Pass scalar"
" int values and vmap over the batch dimension."
)
slices[self.domain.axes((attr,))[0]] = val
values = self.values[tuple(slices)]
domain = self.domain.marginalize(relevant)
return Factor(domain, values)
[docs]
def supports(self, attrs: Attribute | Sequence[Attribute]) -> bool:
return self.domain.supports(attrs)
# Functions that operate element-wise
[docs]
def exp(self) -> Factor:
"""Applies element-wise exponentiation (jnp.exp) to the factor's values."""
return Factor(self.domain, jnp.exp(self.values))
[docs]
def log(self) -> Factor:
"""Applies element-wise logarithm (jnp.log) to the factor's values."""
return Factor(self.domain, jnp.log(self.values))
[docs]
def normalize(
self, total: jax.Array | float = 1.0, log: bool = False
) -> Factor:
"""Normalizes the factor so its values sum to `total` (or log-normalize)."""
if log:
return self + jnp.log(total) - self.logsumexp()
return self * total / self.sum()
[docs]
def copy(self) -> Factor:
"""Returns a copy of the factor (potentially shallow due to JAX)."""
return self
def __float__(self) -> float:
if len(self.domain) > 0:
raise ValueError("Domain must be empty to convert to float.")
return float(self.values)
# Binary operations between two factors
def _binaryop(self, fn: Callable, other: Factor | ArrayLike) -> Factor:
"""Helper for applying binary operations between this factor and another factor or scalar."""
if not isinstance(other, Factor):
assert jnp.ndim(other) == 0
other = Factor(Domain([], []), jnp.asarray(other))
newdom = self.domain.merge(other.domain)
factor1 = self.expand(newdom)
factor2 = other.expand(newdom)
return Factor(newdom, fn(factor1.values, factor2.values))
def __sub__(self, other: Factor | ArrayLike) -> Factor:
return self._binaryop(jnp.subtract, other)
def __truediv__(self, other: Factor | ArrayLike) -> Factor:
return self._binaryop(jnp.divide, other)
def __mul__(self, other: Factor | ArrayLike) -> Factor:
"""Multiply two factors together.
Example Usage:
>>> f1 = Factor.ones(Domain(['a','b'], [2,3]))
>>> f2 = Factor.ones(Domain(['b','c'], [3,4]))
>>> f3 = f1 * f2
>>> print(f3.domain)
Domain(a: 2, b: 3, c: 4)
Args:
other: the other factor to multiply
Returns:
the product of the two factors
"""
return self._binaryop(jnp.multiply, other)
def __add__(self, other: Factor | ArrayLike) -> Factor:
return self._binaryop(jnp.add, other)
def __radd__(self, other: ArrayLike) -> Factor:
return self + other
def __rsub__(self, other: ArrayLike) -> Factor:
return self + (-1 * other)
def __rmul__(self, other: ArrayLike) -> Factor:
return self * other
def __neg__(self) -> Factor:
return self * -1
[docs]
def dot(self, other: Factor) -> ArrayLike:
if self.domain != other.domain:
raise ValueError(f"Domains do not match {self.domain} != {other.domain}")
return jnp.sum(
self.values * other.values,
where=(self.values != 0) & (other.values != 0),
)
[docs]
def datavector(self, flatten: bool = True) -> jax.Array:
"""Returns the factor's values as a flattened vector or original array."""
return self.values.flatten() if flatten else self.values
[docs]
def pad(
self, mesh: jax.sharding.Mesh | None, pad_value: Literal[0, "-inf"]
) -> Factor:
if mesh is None:
return self
pad_amounts = [0] * len(self.domain)
for i, ax in enumerate(self.domain):
if ax in mesh.axis_names:
size = self.domain[ax]
num_shards = mesh.axis_sizes[mesh.axis_names.index(ax)]
pad_amounts[i] = -size % num_shards
values = jnp.pad(
self.values,
pad_width=tuple((0, w) for w in pad_amounts),
constant_values=0.0 if pad_value == 0 else -jnp.inf,
)
# We keep the domain as-is here, even though values is now larger.
# We have a couple of options
# 1. Explicitly unpad when we are done.
# 2. Never unpad, just expand the domain, and make sure new elements are impossible.
# 3. Allow values to be an array where each dim is >= the domain implies, and truncate
# when necessary.
return Factor(self.domain, values)
[docs]
def apply_sharding(self, mesh: jax.sharding.Mesh | None) -> Factor:
"""Apply sharding constraint to the factor values.
The sharding strategy is automatically determined based on the provided
mesh, and the factor domain.
Args:
mesh: The mesh over which the factor should be sharded.
Returns:
A new factor identical to self with sharding constraints applied to the values.
"""
if mesh is None:
return self
pspec: list[Attribute | None] = [None] * len(self.domain)
for i, ax in enumerate(self.domain):
if ax in mesh.axis_names:
pspec[i] = ax
sharding = jax.sharding.NamedSharding(
mesh, jax.sharding.PartitionSpec(*pspec)
)
return Factor(
domain=self.domain,
values=jax.lax.with_sharding_constraint(self.values, sharding),
)
# Re-export for backward compatibility; canonical definition is in _api.py.
from ._api import Projectable # noqa: E402 # pylint: disable=wrong-import-position,unused-import
__all__ = [name for name in dir() if not name.startswith("_")] + [
"Projectable",
]