"""Source-generated *fused* numba kernels for composite / mixture trees.
The generated-numba kernels in :mod:`declarations` lower one leaf at a time: each leaf materializes its
sufficient statistic in numpy and runs its own row loop, so a composite or mixture pays a Python<->C
boundary crossing and an intermediate allocation *per factor and per component*. For a deep model of
low-cost leaves that overhead dominates -- numpy itself is multi-pass for the same reason.
This module instead emits a SINGLE ``@njit`` function for the whole model structure: one pass over the
rows with the composite sum and the mixture log-sum-exp in nopython registers, no per-factor allocation.
Two leaf flavours, each using the right primitive (numba supports BOTH a SIMD-vectorized scalar loop and
BLAS via ``@``/``np.dot`` in nopython):
* **scalar leaves** (Gaussian, Exponential, Poisson, Gamma, Geometric, Bernoulli): univariate, O(1) per
element. Their scalar log-density is inlined into the row loop -- numba SIMD-vectorizes it, and fusing
every factor into one pass with no intermediate arrays beats numpy's multi-pass evaluation. A leaf may
have ``arity > 1`` data arrays (e.g. Poisson carries ``(x, lgamma(x+1))``, Gamma ``(x, log x)``).
* **matrix leaves** (MultivariateGaussian, ...): the quadratic form ``(x-mu)' P (x-mu)`` and the E-step's
weighted Gram ``X' (r .* X)`` are matmuls -- emitted as numba ``@`` (BLAS). Fusing the score, the
responsibility soft-max and the BLAS accumulation into one njit beats numpy's per-component dispatch
(~5x on a GMM E-step). A scalar triple-loop here would be ~15x SLOWER -- never hand-roll a matmul.
When a model mixes the two, the matrix quad-forms are precomputed via BLAS before the row loop, the row
loop scores + forms responsibilities (storing them only when a matrix leaf needs them downstream), and the
matrix sufficient statistics are accumulated via BLAS after it. Pure-scalar models never materialize the
responsibility matrix, so their fast path is byte-for-byte the original one.
A discrete leaf whose value-keyed category counts are the sufficient statistic (Categorical) fuses too --
the **categorical** kind scores from a precomputed (K, C) log-prob table and accumulates a (K, C) weighted
count histogram. The general bar for fusion is a *map-reducible* sufficient statistic (a sum/count, a
min/max, or a histogram -- not necessarily additive): that is what lets Binomial (n from max x) and
NegativeBinomial (iterative dispersion over the count histogram) fuse despite non-additive/iterative MLEs.
Beyond the templated kinds, coverage is COMPLETE for protocol-bearing factors:
* **proper chain laws** (MarkovChain with an explicit length model): handled by the native bridge so
their length probability and versioned sufficient-statistic contract remain intact.
* **bridged factors** (everything else inside a composite -- nested Mixtures, HMMs, nested Composites,
length-model chains, untemplated leaves like Laplace): scored from a precomputed per-component table
(each column the factor's own native ``seq_log_density``), estimated by the factor's own accumulator
driven by the responsibility matrix -- native semantics, priors and all, while the softmax and every
templated sibling stay in one nopython pass. Registered last; specific templates always outrank it.
Every scorer and E-step has a chunk-parallel prange variant (fixed chunking, fixed-order combine:
bit-stable across reruns and worker counts), honors reduced-precision compute (float32 rows, float64
accumulation), and the scorers optionally take the qlut quantized log-sum-exp (``lse_bits``; error
bounded by the grid half-step, compounding per mixture level on nested trees). The nested scalar-tree
kernels (:mod:`fused_nested`) carry the same contracts.
Principled exclusions, deliberate and documented rather than pending: E-steps always use EXACT exp (a
delta-grid perturbation could flip a monotone-gate accept); bare nested mixtures keep fused_nested's
in-kernel path (faster than bridging), and only when fused_nested itself declines (e.g. minmax leaves
like Pareto) does a LAST-RESORT bare-bridge pass cover the shape -- host-exact semantics, one softmax
kernel; GradLeaf M-steps stay eager torch (torch.compile measured 0.79-0.93x on CPU -- see the GradLeaf
docstring).
Generated kernels are compiled once and disk-cached (see :func:`_njit`), so the compile cost is paid once
per structure *ever*, not per process; :func:`clean_fused_cache` garbage-collects entries stranded by
structure changes or digest-salt bumps.
"""
from __future__ import annotations
import hashlib
import marshal
from collections.abc import Callable
from dataclasses import dataclass, field
from typing import Any
import numpy as np
from mixle.stats.compute.mixture_evidence import (
InvalidMixtureEvidenceError,
raise_for_invalid_log_evidence,
)
# --- leaf templates -------------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class LeafTemplate:
"""How to fuse one leaf family.
``kind="scalar"`` leaves use ``expr`` / ``acc_stmt`` over ``arity`` 1-D data arrays (the row values
are passed as a list ``vals`` of numba expressions). ``kind="matrix"`` leaves use the ``mat_*`` hooks
(a single 2-D data array + BLAS precompute/accumulate). ``params`` stacks the leaf's parameters across
the K mixture components; ``to_value`` packs a component's accumulated statistics into the estimator's
``value()`` tuple.
"""
name: str
matches: Callable[[Any], bool]
data: Callable[[Any], tuple[np.ndarray, ...]] # leaf encoding -> tuple of (N,...) data arrays
params: Callable[[list[Any]], dict[str, np.ndarray]] # K leaf dists -> {pname: (K,...) array}
to_value: Callable[[tuple, float], tuple] | None = None # (stats for comp k, count_k) -> leaf value
arity: int = 1 # number of data arrays this leaf consumes
kind: str = "scalar"
# scalar hooks (vals = [numba expr per data array], p/a = {name: arg}, r = responsibility var)
expr: Callable[[list[str], dict[str, str]], str] | None = None
acc_names: tuple[str, ...] = () # per-leaf weighted-statistic accumulator arrays, shape (K,)
acc_stmt: Callable[[list[str], dict[str, str], str], str] | None = None
# matrix hooks (emit numba source lines; leaf index `i`, {pname: arg} map)
mat_precompute: Callable[[int, dict[str, str]], list[str]] | None = None # -> fill q{i} (K,N)
mat_row: Callable[[int, dict[str, str]], str] | None = None # -> per-(i,k) scoring expr
mat_accumulate: Callable[[int], list[str]] | None = None # -> fill S1_{i} (K,D), S2_{i} (K,D,D) via BLAS
# vector hooks (2-D data, inline per-dim loop, (K,D) accumulators; emit source lines)
vec_row: Callable[[int, dict[str, str]], list[str]] | None = None # -> add the leaf score to `acc`
vec_accumulate: Callable[[int, dict[str, str]], list[str]] | None = None # -> per-dim weighted-stat update
# tabulated hooks (low-cardinality integer x: precompute a (K, max+1) log-pmf table, look it up per row;
# accumulate sum_wx and -- when ``tab_hist`` -- a per-component weighted histogram (a map-reducible
# bincount). This is how discrete families whose MLE needs a max/min reduction (Binomial) or the full
# count distribution (NegativeBinomial's iterative dispersion) still fuse in one pass.)
tab_table: Callable[[list[Any], int], np.ndarray] | None = None # (comps, max_x) -> (K, max_x+1) log-pmf
tab_hist: bool = False # also accumulate a (K, max_x+1) weighted histogram
tab_to_value: Callable[[float, float, np.ndarray | None, int, int], tuple] | None = (
None # (sx,count,hist_k,min,max)
)
# global min/max over x (a map-reducible reduction) for a scalar leaf whose value() needs it (Pareto's xm)
wants_minmax: bool = False
to_value_g: Callable[[tuple, float, float, float], tuple] | None = None # (stats, count, min, max) -> value
# scalar accumulators combined by MIN instead of SUM (allocated +inf; the parallel chunk axis and any
# merge must min-reduce them) -- per-component weighted support minima, e.g. Pareto's xm
min_acc_names: tuple[str, ...] = ()
# categorical hooks (data is already a category index; score from a (K,C) log-prob table, accumulate a
# (K,C) weighted count histogram -- the only sufficient statistic). Generalizes the tabulated pattern to
# value-keyed categories (the table/value depend on the encoding's category list, so they get ``enc``).
cat_table: Callable[[list[Any], Any], np.ndarray] | None = None # (comps, enc) -> (K, C) log-prob table
cat_to_value: Callable[[np.ndarray, Any, float], Any] | None = None # (hist_k, enc, count) -> leaf value
# chain hooks (a Markov-chain factor: RAGGED per-row sequences, already encoded as scatter-ready
# (row_idx, state_idx) arrays). Scoring precomputes a per-row per-component score table by
# scatter-adding init/transition log-probabilities BEFORE the row loop (the same precompute slot the
# matrix kind's BLAS quad forms use), the row fragment is a table lookup, and the E-step accumulates
# weighted init/transition histograms from the responsibility matrix R in the post pass.
chain_tables: Callable[[list[Any], Any], tuple[np.ndarray, np.ndarray]] | None = (
None # (comps, enc) -> (initT (K,S), transT (K,S,S)) in the ENCODER's state indexing
)
chain_to_value: Callable[[np.ndarray, np.ndarray, Any, float], Any] | None = (
None # (init_hist_k (S,), trans_hist_k (S,S), enc, count) -> leaf value
)
dtype: str = "float64"
cache_version: str = "1"
fingerprint: str = field(init=False, repr=False, compare=True)
def __post_init__(self) -> None:
if not isinstance(self.name, str) or not self.name:
raise ValueError("leaf template name must be a non-empty string.")
if not isinstance(self.cache_version, str) or not self.cache_version:
raise ValueError("leaf template cache_version must be a non-empty string.")
if isinstance(self.arity, bool) or not isinstance(self.arity, int) or self.arity < 0:
raise ValueError("leaf template arity must be a non-negative integer.")
if not callable(self.matches) or not callable(self.data) or not callable(self.params):
raise TypeError("leaf template matches, data, and params hooks must be callable.")
object.__setattr__(self, "fingerprint", _leaf_template_fingerprint(self))
_TEMPLATES: list[LeafTemplate] = []
def _stable_semantic_token(value: Any) -> Any:
"""Return a deterministic token for Python hook code and captured state."""
if value is None or isinstance(value, (bool, int, float, str, bytes)):
return (type(value).__name__, value)
if isinstance(value, (tuple, list)):
return (type(value).__name__, tuple(_stable_semantic_token(item) for item in value))
if isinstance(value, dict):
items = sorted(
((_stable_semantic_token(key), _stable_semantic_token(item)) for key, item in value.items()),
key=repr,
)
return ("dict", tuple(items))
if isinstance(value, np.ndarray):
return (
"ndarray",
value.dtype.str,
tuple(value.shape),
hashlib.sha256(np.ascontiguousarray(value).tobytes()).hexdigest(),
)
code = getattr(value, "__code__", None)
if code is not None:
closure = getattr(value, "__closure__", None) or ()
return (
"python-callable",
getattr(value, "__module__", None),
getattr(value, "__qualname__", None),
hashlib.sha256(marshal.dumps(code)).hexdigest(),
_stable_semantic_token(getattr(value, "__defaults__", None)),
_stable_semantic_token(getattr(value, "__kwdefaults__", None)),
tuple(_stable_semantic_token(cell.cell_contents) for cell in closure),
)
explicit = getattr(value, "__mixle_cache_key__", None)
if explicit is not None:
return ("explicit", _stable_semantic_token(explicit))
state = getattr(value, "__dict__", None)
return (
"opaque",
type(value).__module__,
type(value).__qualname__,
_stable_semantic_token(state) if state else None,
)
def _leaf_template_fingerprint(template: LeafTemplate) -> str:
"""Fingerprint every field that can change generated or marshalled semantics."""
hooks = (
template.matches,
template.data,
template.params,
template.to_value,
template.expr,
template.acc_stmt,
template.mat_precompute,
template.mat_row,
template.mat_accumulate,
template.vec_row,
template.vec_accumulate,
template.tab_table,
template.tab_to_value,
template.to_value_g,
template.cat_table,
template.cat_to_value,
template.chain_tables,
template.chain_to_value,
)
semantics = (
template.name,
template.cache_version,
template.arity,
template.kind,
template.acc_names,
template.tab_hist,
template.wants_minmax,
template.min_acc_names,
template.dtype,
tuple(_stable_semantic_token(hook) for hook in hooks),
)
return hashlib.sha256(repr(semantics).encode("utf-8")).hexdigest()
[docs]
def register_leaf_template(t: LeafTemplate) -> None:
"""Register a uniquely named, fingerprinted leaf template.
Extensions registered after the built-in bridge are inserted before it so
the catch-all bridge cannot shadow them. Replacing a registered name is
forbidden: create a new process with the changed template (and bump
``cache_version`` when semantics live outside inspectable Python hook code).
"""
if not isinstance(t, LeafTemplate):
raise TypeError("register_leaf_template requires a LeafTemplate.")
if any(existing.name == t.name for existing in _TEMPLATES):
raise ValueError("leaf template name %r is already registered." % t.name)
bridge_index = next((index for index, existing in enumerate(_TEMPLATES) if existing.kind == "bridge"), None)
if bridge_index is None or t.kind == "bridge":
_TEMPLATES.append(t)
else:
_TEMPLATES.insert(bridge_index, t)
def _template_for(dist: Any, allow_bridge: bool = True) -> LeafTemplate | None:
for t in _TEMPLATES:
if (allow_bridge or t.kind != "bridge") and t.matches(dist):
return t
return None
def _arr(enc: Any, dtype: Any = np.float64) -> np.ndarray:
return np.ascontiguousarray(enc, dtype=dtype)
def _on_support(expression: str, condition: str) -> str:
"""Guard an inlined log-density expression with its encoded support."""
return f"(({expression}) if {condition} else -np.inf)"
# --- scalar leaves --------------------------------------------------------------------------------
def _gaussian_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.array([c.mu for c in comps], dtype=np.float64)
s2 = np.array([c.sigma2 for c in comps], dtype=np.float64)
return {"mu": mu, "inv2s2": 0.5 / s2, "lognorm": -0.5 * np.log(2.0 * np.pi * s2)}
register_leaf_template(
LeafTemplate(
name="gaussian",
matches=lambda d: type(d).__name__ == "GaussianDistribution",
data=lambda enc: (_arr(enc),),
params=_gaussian_params,
expr=lambda v, p: f"{p['lognorm']}[k] - ({v[0]} - {p['mu']}[k]) * ({v[0]} - {p['mu']}[k]) * {p['inv2s2']}[k]",
acc_names=("sx", "sx2"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (s[0], s[1], count, count), # (sum_wx, sum_wx2, sum_w, sum_w)
)
)
register_leaf_template(
LeafTemplate(
name="exponential",
matches=lambda d: type(d).__name__ == "ExponentialDistribution",
data=lambda enc: (_arr(enc),),
params=lambda comps: (lambda beta: {"rate": 1.0 / beta, "lograte": -np.log(beta)})(
np.array([c.beta for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
f"{p['lograte']}[k] - {p['rate']}[k] * {v[0]}",
f"np.isfinite({v[0]}) and {v[0]} >= 0.0",
),
acc_names=("sx",),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (sum_w, sum_wx)
)
)
register_leaf_template(
LeafTemplate(
name="geometric",
matches=lambda d: type(d).__name__ == "GeometricDistribution",
data=lambda enc: (_arr(enc),), # support {1,2,...}: log p((x)) = (x-1) ln(1-p) + ln p
params=lambda comps: (lambda p: {"logp": np.log(p), "log1mp": np.log1p(-p)})(
np.array([c.p for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
(
f"(0.0 if {p['log1mp']}[k] == -np.inf and {v[0]} == 1.0 "
f"else {p['logp']}[k] + ({v[0]} - 1.0) * {p['log1mp']}[k])"
),
f"np.isfinite({v[0]}) and {v[0]} >= 1.0 and np.floor({v[0]}) == {v[0]}",
),
acc_names=("sx",),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (sum_w, sum_wx)
)
)
register_leaf_template(
LeafTemplate(
name="bernoulli",
matches=lambda d: type(d).__name__ == "BernoulliDistribution",
data=lambda enc: (_arr(enc),), # log p(x) = x*logit + ln(1-p)
params=lambda comps: (lambda p: {"logp": np.log(p), "log1mp": np.log1p(-p)})(
np.array([c.p for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
f"({p['logp']}[k] if {v[0]} == 1.0 else {p['log1mp']}[k])",
f"np.isfinite({v[0]}) and ({v[0]} == 0.0 or {v[0]} == 1.0)",
),
acc_names=("sx",),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (sum_w, sum_wx)
)
)
register_leaf_template(
LeafTemplate(
name="poisson",
matches=lambda d: type(d).__name__ == "PoissonDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (x, lgamma(x+1)); log p = -lam + x ln lam - lgamma(x+1)
params=lambda comps: (lambda lam: {"lam": lam, "loglam": np.log(lam)})(
np.array([c.lam for c in comps], dtype=np.float64)
),
arity=2,
expr=lambda v, p: _on_support(
f"{p['loglam']}[k] * {v[0]} - {p['lam']}[k] - {v[1]}",
f"np.isfinite({v[0]}) and {v[0]} >= 0.0 and np.floor({v[0]}) == {v[0]}",
),
acc_names=("sx",),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (sum_w, sum_wx)
)
)
def _gamma_params(comps: list[Any]) -> dict[str, np.ndarray]:
k = np.array([c.k for c in comps], dtype=np.float64)
theta = np.array([c.theta for c in comps], dtype=np.float64)
from scipy.special import gammaln
return {"km1": k - 1.0, "inv_theta": 1.0 / theta, "norm": -k * np.log(theta) - gammaln(k)}
register_leaf_template(
LeafTemplate(
name="gamma",
matches=lambda d: type(d).__name__ == "GammaDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (x, log x); log p = (k-1)lnx - x/th - k ln th - lgamma(k)
params=_gamma_params,
arity=2,
expr=lambda v, p: _on_support(
f"{p['norm']}[k] + {p['km1']}[k] * {v[1]} - {p['inv_theta']}[k] * {v[0]}",
f"np.isfinite({v[0]}) and {v[0]} > 0.0",
),
acc_names=("sx", "slogx"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['slogx']}[k] += {r} * {v[1]}",
to_value=lambda s, count: (count, s[0], s[1]), # (sum_w, sum_wx, sum_w_logx)
)
)
def _arr0(enc: Any) -> tuple[np.ndarray]:
"""Raw x of a univariate encoding (enc[0] for the precomputed-transform tuples, else the array)."""
return (_arr(enc[0] if isinstance(enc, tuple) else enc),)
# Each of these scalar leaves takes raw x and computes its transforms inline (numba has log/log1p/cos/sin),
# so only the per-component normalizer is precomputed in numpy. Verified bit-accurate vs seq_log_density.
register_leaf_template(
LeafTemplate(
name="halfnormal",
matches=lambda d: type(d).__name__ == "HalfNormalDistribution",
data=_arr0,
params=lambda comps: (lambda s2: {"a": 0.5 / s2, "lognorm": 0.5 * np.log(2.0 / np.pi) - 0.5 * np.log(s2)})(
np.array([c.sigma**2 for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
f"{p['lognorm']}[k] - {v[0]} * {v[0]} * {p['a']}[k]",
f"np.isfinite({v[0]}) and {v[0]} >= 0.0",
),
acc_names=("sx2",),
acc_stmt=lambda v, a, r: f"{a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (n, sum_w x^2)
)
)
register_leaf_template(
LeafTemplate(
name="rayleigh",
matches=lambda d: type(d).__name__ == "RayleighDistribution",
data=_arr0,
params=lambda comps: (lambda s2: {"a": 0.5 / s2, "lognorm": -np.log(s2)})(
np.array([c.sigma**2 for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
f"{p['lognorm']}[k] + np.log({v[0]}) - {v[0]} * {v[0]} * {p['a']}[k]",
f"np.isfinite({v[0]}) and {v[0]} > 0.0",
),
acc_names=("sx2",),
acc_stmt=lambda v, a, r: f"{a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (n, sum_w x^2)
)
)
def _inverse_gaussian_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.array([c.mu for c in comps], dtype=np.float64)
lam = np.array([c.lam for c in comps], dtype=np.float64)
return {
"c1": -lam / (2.0 * mu * mu),
"c2": -lam / 2.0,
"lognorm": 0.5 * (np.log(lam) - np.log(2.0 * np.pi)) + lam / mu,
}
register_leaf_template(
LeafTemplate(
name="inversegaussian",
matches=lambda d: type(d).__name__ == "InverseGaussianDistribution",
data=_arr0,
params=_inverse_gaussian_params,
expr=lambda v, p: _on_support(
(f"{p['lognorm']}[k] - 1.5 * np.log({v[0]}) + {p['c1']}[k] * {v[0]} + {p['c2']}[k] / {v[0]}"),
f"np.isfinite({v[0]}) and {v[0]} > 0.0",
),
acc_names=("sx", "s1x"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['s1x']}[k] += {r} / {v[0]}",
to_value=lambda s, count: (count, s[0], s[1]), # (n, sum_w x, sum_w / x)
)
)
def _beta_params(comps: list[Any]) -> dict[str, np.ndarray]:
from scipy.special import betaln
a = np.array([c.a for c in comps], dtype=np.float64)
b = np.array([c.b for c in comps], dtype=np.float64)
return {"am1": a - 1.0, "bm1": b - 1.0, "lognorm": -betaln(a, b)}
register_leaf_template(
LeafTemplate(
name="beta",
matches=lambda d: type(d).__name__ == "BetaDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1]), _arr(enc[2]), _arr(enc[3])), # (log x, log(1-x), x, x^2)
params=_beta_params,
arity=4,
expr=lambda v, p: _on_support(
f"{p['lognorm']}[k] + {p['am1']}[k] * {v[0]} + {p['bm1']}[k] * {v[1]}",
f"np.isfinite({v[2]}) and {v[2]} > 0.0 and {v[2]} < 1.0",
),
acc_names=("slogx", "slog1mx", "sx", "sx2"),
acc_stmt=lambda v, a, r: (
f"{a['slogx']}[k] += {r} * {v[0]}; {a['slog1mx']}[k] += {r} * {v[1]}; "
f"{a['sx']}[k] += {r} * {v[2]}; {a['sx2']}[k] += {r} * {v[3]}"
),
to_value=lambda s, count: (count, s[0], s[1], s[2], s[3]), # (n, sumlogx, sumlog1mx, sumx, sumx2)
)
)
def _inverse_gamma_params(comps: list[Any]) -> dict[str, np.ndarray]:
from scipy.special import gammaln
a = np.array([c.alpha for c in comps], dtype=np.float64)
b = np.array([c.beta for c in comps], dtype=np.float64)
return {"ap1": a + 1.0, "beta": b, "lognorm": a * np.log(b) - gammaln(a)}
register_leaf_template(
LeafTemplate(
name="inversegamma",
matches=lambda d: type(d).__name__ == "InverseGammaDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (log x, 1/x)
params=_inverse_gamma_params,
arity=2,
expr=lambda v, p: _on_support(
f"{p['lognorm']}[k] - {p['ap1']}[k] * {v[0]} - {p['beta']}[k] * {v[1]}",
f"np.isfinite({v[0]}) and np.isfinite({v[1]}) and {v[1]} > 0.0",
),
acc_names=("s1x", "slogx"),
acc_stmt=lambda v, a, r: f"{a['s1x']}[k] += {r} * {v[1]}; {a['slogx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0], -s[1]), # (n, sum_w / x, sum_w log(1/x)) -- Gamma on 1/x
)
)
def _vonmises_params(comps: list[Any]) -> dict[str, np.ndarray]:
from scipy.special import i0
mu = np.array([c.mu for c in comps], dtype=np.float64)
kappa = np.array([c.kappa for c in comps], dtype=np.float64)
return {"kcos": kappa * np.cos(mu), "ksin": kappa * np.sin(mu), "lognorm": -np.log(2.0 * np.pi * i0(kappa))}
register_leaf_template(
LeafTemplate(
name="vonmises",
matches=lambda d: type(d).__name__ == "VonMisesDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (cos x, sin x)
params=_vonmises_params,
arity=2,
expr=lambda v, p: f"{p['lognorm']}[k] + {p['kcos']}[k] * {v[0]} + {p['ksin']}[k] * {v[1]}",
acc_names=("scos", "ssin"),
acc_stmt=lambda v, a, r: f"{a['scos']}[k] += {r} * {v[0]}; {a['ssin']}[k] += {r} * {v[1]}",
to_value=lambda s, count: (count, s[0], s[1]), # (n, sum_w cos, sum_w sin)
)
)
def _wrappedcauchy_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.array([c.mu for c in comps], dtype=np.float64)
rho = np.array([c.rho for c in comps], dtype=np.float64)
return {
"a": 1.0 + rho * rho,
"rcos": 2.0 * rho * np.cos(mu),
"rsin": 2.0 * rho * np.sin(mu),
"lognorm": np.log1p(-rho * rho) - np.log(2.0 * np.pi),
}
register_leaf_template(
LeafTemplate(
name="wrappedcauchy",
matches=lambda d: type(d).__name__ == "WrappedCauchyDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (cos x, sin x)
params=_wrappedcauchy_params,
arity=2,
expr=lambda v, p: (
f"{p['lognorm']}[k] - np.log({p['a']}[k] - {p['rcos']}[k] * {v[0]} - {p['rsin']}[k] * {v[1]})"
),
acc_names=("scos", "ssin"),
acc_stmt=lambda v, a, r: f"{a['scos']}[k] += {r} * {v[0]}; {a['ssin']}[k] += {r} * {v[1]}",
to_value=lambda s, count: (s[0], s[1], count), # (sum_w cos, sum_w sin, n)
)
)
register_leaf_template(
LeafTemplate(
name="pareto",
matches=lambda d: type(d).__name__ == "ParetoDistribution",
data=lambda enc: (_arr(enc[0]), _arr(enc[1])), # (x, log x)
params=lambda comps: (lambda al, xm: {"const": np.log(al) + al * np.log(xm), "am1": al + 1.0, "xm": xm})(
np.array([c.alpha for c in comps], dtype=np.float64), np.array([c.xm for c in comps], dtype=np.float64)
),
arity=2,
# support x >= xm: log p = log a + a log xm - (a+1) log x; below the scale the component is impossible
expr=lambda v, p: f"(({p['const']}[k] - {p['am1']}[k] * {v[1]}) if {v[0]} >= {p['xm']}[k] else -np.inf)",
acc_names=("slogx", "wmin"),
# wmin tracks the PER-COMPONENT minimum over rows with responsibility > 0 -- the legacy
# ParetoAccumulator statistic that ParetoEstimator uses directly as the scale xm. The global
# data minimum (to_value_g's mn) is wrong whenever component supports differ.
acc_stmt=lambda v, a, r: (
f"{a['slogx']}[k] += {r} * {v[1]}; "
f"{a['wmin']}[k] = {v[0]} if {r} > 0.0 and {v[0]} < {a['wmin']}[k] else {a['wmin']}[k]"
),
wants_minmax=True, # keeps minmax leaves out of the nested emitter (no to_value / min plumbing there)
min_acc_names=("wmin",),
to_value_g=lambda s, count, mn, mx: (count, s[0], s[1]), # (n, sum_w logx, per-component min x)
)
)
register_leaf_template(
LeafTemplate(
name="logseries",
matches=lambda d: type(d).__name__ == "LogSeriesDistribution",
data=_arr0,
params=lambda comps: (lambda p: {"logp": np.log(p), "lognorm": -np.log(-np.log1p(-p))})(
np.array([c.p for c in comps], dtype=np.float64)
),
expr=lambda v, p: _on_support(
f"{p['lognorm']}[k] + {v[0]} * {p['logp']}[k] - np.log({v[0]})",
f"np.isfinite({v[0]}) and {v[0]} >= 1.0 and np.floor({v[0]}) == {v[0]}",
),
acc_names=("sx",),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}",
to_value=lambda s, count: (count, s[0]), # (n, sum_w x)
)
)
# Location-scale families with moment-matched estimators: value() = (sum_w x, sum_w x^2, n), exactly the
# Gaussian statistics; only the (closed-form) density differs. z = (x - loc) / scale, inlined.
def _locscale_params(comps: list[Any]) -> dict[str, np.ndarray]:
loc = np.array([c.loc for c in comps], dtype=np.float64)
scale = np.array([c.scale for c in comps], dtype=np.float64)
return {"loc": loc, "inv_scale": 1.0 / scale, "nls": -np.log(scale)}
def _z(v: list[str], p: dict[str, str]) -> str:
return f"(({v[0]} - {p['loc']}[k]) * {p['inv_scale']}[k])"
register_leaf_template(
LeafTemplate(
name="gumbel",
matches=lambda d: type(d).__name__ == "GumbelDistribution",
data=lambda enc: (_arr(enc),),
params=_locscale_params,
expr=lambda v, p: f"{p['nls']}[k] - {_z(v, p)} - np.exp(-{_z(v, p)})",
acc_names=("sx", "sx2"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (s[0], s[1], count), # (sum_w x, sum_w x^2, n)
)
)
register_leaf_template(
LeafTemplate(
name="logistic",
matches=lambda d: type(d).__name__ == "LogisticDistribution",
data=lambda enc: (_arr(enc),),
params=_locscale_params,
expr=lambda v, p: f"{p['nls']}[k] - {_z(v, p)} - 2.0 * np.log1p(np.exp(-{_z(v, p)}))",
acc_names=("sx", "sx2"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (s[0], s[1], count), # (sum_w x, sum_w x^2, n)
)
)
def _studentt_params(comps: list[Any]) -> dict[str, np.ndarray]:
from scipy.special import gammaln
loc = np.array([c.loc for c in comps], dtype=np.float64)
scale = np.array([c.scale for c in comps], dtype=np.float64)
df = np.array([c.df for c in comps], dtype=np.float64)
lognorm = gammaln(0.5 * (df + 1.0)) - gammaln(0.5 * df) - 0.5 * np.log(df * np.pi) - np.log(scale)
return {"loc": loc, "inv_scale": 1.0 / scale, "ndf": 0.5 * (df + 1.0), "inv_df": 1.0 / df, "lognorm": lognorm}
register_leaf_template(
LeafTemplate(
name="studentt",
matches=lambda d: type(d).__name__ == "StudentTDistribution",
data=lambda enc: (_arr(enc),),
params=_studentt_params,
expr=lambda v, p: f"{p['lognorm']}[k] - {p['ndf']}[k] * np.log1p({_z(v, p)} * {_z(v, p)} * {p['inv_df']}[k])",
acc_names=("sx", "sx2"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (s[0], s[1], count), # (sum_w x, sum_w x^2, n) -- moment-matched loc/scale
)
)
def _loggaussian_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.array([c.mu for c in comps], dtype=np.float64)
s2 = np.array([c.sigma2 for c in comps], dtype=np.float64)
return {"mu": mu, "inv2s2": 0.5 / s2, "lognorm": -0.5 * np.log(2.0 * np.pi * s2)}
register_leaf_template(
LeafTemplate(
name="loggaussian",
matches=lambda d: type(d).__name__ == "LogGaussianDistribution",
data=lambda enc: (_arr(enc),), # the encoder already supplies log x; score = Gaussian(log x) - log x (Jacobian)
params=_loggaussian_params,
expr=lambda v, p: (
f"{p['lognorm']}[k] - {v[0]} - ({v[0]} - {p['mu']}[k]) * ({v[0]} - {p['mu']}[k]) * {p['inv2s2']}[k]"
),
acc_names=("sx", "sx2"),
acc_stmt=lambda v, a, r: f"{a['sx']}[k] += {r} * {v[0]}; {a['sx2']}[k] += {r} * {v[0]} * {v[0]}",
to_value=lambda s, count: (s[0], s[1], count, count), # (sum_w_logx, sum_w_logx2, sum_w, sum_w)
)
)
# --- vector leaves (2-D data, per-dim inline loop, (K,D) accumulators -- no matmul) ----------------
def _diaggaussian_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.ascontiguousarray(np.stack([np.asarray(c.mu, dtype=np.float64) for c in comps]))
s2 = np.ascontiguousarray(np.stack([np.asarray(c.covar, dtype=np.float64) for c in comps])) # diagonal variances
lognorm = -0.5 * np.sum(np.log(2.0 * np.pi * s2), axis=1) # (K,)
return {"mu": mu, "inv2s2": 0.5 / s2, "lognorm": lognorm}
def _diag_row(i: int, p: dict[str, str]) -> list[str]:
return [
f"acc += {p['lognorm']}[k]",
f"for d in range(x{i}_0.shape[1]):",
f" diff{i} = x{i}_0[i, d] - {p['mu']}[k, d]",
f" acc -= diff{i} * diff{i} * {p['inv2s2']}[k, d]",
]
def _diag_accumulate(i: int, a: dict[str, str]) -> list[str]:
return [
f"for d in range(x{i}_0.shape[1]):",
f" {a['s1']}[k, d] += r * x{i}_0[i, d]",
f" {a['s2']}[k, d] += r * x{i}_0[i, d] * x{i}_0[i, d]",
]
register_leaf_template(
LeafTemplate(
name="diaggaussian",
matches=lambda d: type(d).__name__ == "DiagonalGaussianDistribution",
data=lambda enc: (_arr(enc),),
params=_diaggaussian_params,
kind="vector",
vec_row=_diag_row,
acc_names=("s1", "s2"),
vec_accumulate=_diag_accumulate,
to_value=lambda s, count: (s[0], s[1], count), # (sum_wx (D,), sum_wxx (D,), sum_w)
)
)
# --- matrix leaves --------------------------------------------------------------------------------
def _mvgaussian_params(comps: list[Any]) -> dict[str, np.ndarray]:
mu = np.ascontiguousarray(np.stack([np.asarray(c.mu, dtype=np.float64) for c in comps]))
prec = np.ascontiguousarray(np.stack([np.asarray(c.inv_covar, dtype=np.float64) for c in comps]))
d = mu.shape[1]
lognorm = np.empty(len(comps), dtype=np.float64)
for k, c in enumerate(comps):
_, logdet = np.linalg.slogdet(np.asarray(c.covar, dtype=np.float64))
lognorm[k] = -0.5 * (d * np.log(2.0 * np.pi) + logdet)
return {"mu": mu, "prec": prec, "lognorm": lognorm}
def _mv_precompute(i: int, p: dict[str, str]) -> list[str]:
# q{i}[k, r] = (x_r - mu_k)' P_k (x_r - mu_k), the centered quadratic form, via BLAS gemm `C @ P`.
return [
f"q{i} = np.empty((kc, n))",
"for k in range(kc):",
f" C{i} = x{i}_0 - {p['mu']}[k]",
f" Z{i} = C{i} @ {p['prec']}[k]",
" for ii in range(n):",
f" qq{i} = 0.0",
f" for jj in range(x{i}_0.shape[1]):",
f" qq{i} += Z{i}[ii, jj] * C{i}[ii, jj]",
f" q{i}[k, ii] = qq{i}",
]
def _mv_accumulate(i: int) -> list[str]:
# MVGaussian weighted stats: S1_k = X' r_k (weighted sum), S2_k = X' (r_k .* X) (weighted Gram), via BLAS.
return [
"for k in range(kc):",
f" Rk{i} = R[:, k].copy()",
f" S1_{i}[k] = x{i}_0.T @ Rk{i}",
f" S2_{i}[k] = x{i}_0.T @ (Rk{i}.reshape(-1, 1) * x{i}_0)",
]
register_leaf_template(
LeafTemplate(
name="mvgaussian",
matches=lambda d: type(d).__name__ == "MultivariateGaussianDistribution",
data=lambda enc: (_arr(enc),),
params=_mvgaussian_params,
kind="matrix",
mat_precompute=_mv_precompute,
mat_row=lambda i, p: f"-0.5 * q{i}[k, i] + {p['lognorm']}[k]",
mat_accumulate=_mv_accumulate,
to_value=lambda s, count: (s[0], s[1], count), # (sum_wx (D,), sum_wxx (D,D), sum_w)
)
)
# --- tabulated leaves (low-cardinality integer x: table-scored + map-reducible histogram/min-max) --
def _binomial_table(comps: list[Any], max_x: int) -> np.ndarray:
from scipy.special import gammaln
xs = np.arange(max_x + 1, dtype=np.float64)
tab = np.full((len(comps), max_x + 1), -np.inf, dtype=np.float64)
for k, c in enumerate(comps):
n = float(c.n)
mv = float(getattr(c, "min_val", 0) or 0)
xx = xs - mv
valid = (xx >= 0) & (xx <= n)
lp = np.log(c.p) if c.p > 0 else -np.inf
l1p = np.log1p(-c.p) if c.p < 1 else -np.inf
with np.errstate(invalid="ignore", divide="ignore"):
row = gammaln(n + 1) - gammaln(xx + 1) - gammaln(n - xx + 1) + xx * lp + (n - xx) * l1p
tab[k, valid] = row[valid]
return tab
register_leaf_template(
LeafTemplate(
name="binomial",
matches=lambda d: type(d).__name__ == "BinomialDistribution",
data=lambda enc: (_arr(enc[2]),), # enc = (unique, inverse, x, min, max); accumulate over raw x
params=lambda comps: {},
kind="tabulated",
tab_table=_binomial_table,
acc_names=("sx",),
# value() = (count, sum, min_val, max_val); min/max are global reductions over x (weight-independent)
tab_to_value=lambda sx, count, hist, mn, mx: (count, sx, mn, mx),
)
)
def _negbinomial_table(comps: list[Any], max_x: int) -> np.ndarray:
from scipy.special import gammaln
xs = np.arange(max_x + 1, dtype=np.float64)
tab = np.empty((len(comps), max_x + 1), dtype=np.float64)
for k, c in enumerate(comps): # NB(r,p): log p = lgamma(x+r) - lgamma(r) - lgamma(x+1) + r ln p + x ln(1-p)
tab[k] = gammaln(xs + c.r) - gammaln(c.r) - gammaln(xs + 1) + c.r * np.log(c.p) + xs * np.log1p(-c.p)
return tab
def _hist_to_dict(hist_k: np.ndarray) -> dict[int, float]:
nz = np.nonzero(hist_k)[0]
return {int(x): float(hist_k[x]) for x in nz}
register_leaf_template(
LeafTemplate(
name="negbinomial",
matches=lambda d: type(d).__name__ == "NegativeBinomialDistribution",
data=lambda enc: (_arr(enc[0]),), # enc = (x, lgamma(x+1)); accumulate sum_wx + weighted histogram
params=lambda comps: {},
kind="tabulated",
tab_table=_negbinomial_table,
acc_names=("sx", "hist"),
tab_hist=True,
# value() = (count, sum, histogram); the iterative dispersion MLE needs the full weighted count dist
tab_to_value=lambda sx, count, hist, mn, mx: (count, sx, _hist_to_dict(hist)),
)
)
# --- categorical leaves (data is already a category index; (K,C) log-prob table + count histogram) ------
def _categorical_table(comps: list[Any], enc: Any) -> np.ndarray:
values = enc[1] # the shared category list the encoder assigned indices over
tab = np.empty((len(comps), len(values)), dtype=np.float64)
for k, c in enumerate(comps):
tab[k] = [float(c.log_density(v)) for v in values] # -inf where a value is outside this component
return tab
register_leaf_template(
LeafTemplate(
name="categorical",
matches=lambda d: type(d).__name__ == "CategoricalDistribution",
data=lambda enc: (_arr(enc[0]),), # enc = (index, values); index is already 0..C-1
params=lambda comps: {},
kind="categorical",
acc_names=("hist",),
cat_table=_categorical_table,
# value() is the weighted count dict {value: weight}
cat_to_value=lambda hist_k, enc, count: {
enc[1][c]: float(hist_k[c]) for c in range(len(enc[1])) if hist_k[c] != 0.0
},
)
)
def _markov_chain_data(enc: Any) -> tuple[np.ndarray, ...]:
"""The chain encoding's scatter arrays: (init row idx, init state, trans row idx, from, to)."""
_, idx0, idx1, init_x, prev_x, next_x, _, _ = enc
def as_i8(value: Any, path: str) -> np.ndarray:
raw = np.asarray(value)
if raw.ndim != 1:
raise ValueError("%s must be one-dimensional." % path)
try:
numeric = np.asarray(raw, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise TypeError("%s must contain integer indices." % path) from exc
if np.any(~np.isfinite(numeric)) or np.any(np.floor(numeric) != numeric):
raise ValueError("%s must contain finite integer indices." % path)
return np.ascontiguousarray(numeric, dtype=np.int64)
return (
as_i8(idx0, "chain initial row indices"),
as_i8(init_x, "chain initial state indices"),
as_i8(idx1, "chain transition row indices"),
as_i8(prev_x, "chain previous-state indices"),
as_i8(next_x, "chain next-state indices"),
)
def _markov_chain_tables(comps: list[Any], enc: Any) -> tuple[np.ndarray, np.ndarray]:
"""Per-component (K,S) init and (K,S,S) transition log tables in the ENCODER's state order.
Exactly replicates MarkovChainDistribution.seq_log_density's lookup: each component maps the
encoder's states through its own key_map (0 = out-of-support default) into its init/transition
log vectors, with the component's log1p(default) normalizer folded into every entry.
"""
inv_key_map = enc[6]
S = len(inv_key_map)
init_t = np.empty((len(comps), max(S, 1)), dtype=np.float64)
trans_t = np.empty((len(comps), max(S, 1), max(S, 1)), dtype=np.float64)
for k, c in enumerate(comps):
loc = np.asarray([c.key_map.get(u, 0) for u in inv_key_map], dtype=np.int64)
if S:
init_t[k] = c.init_log_pvec[loc] - c.log1p_dv
dense = np.asarray(c.trans_log_pvec.toarray(), dtype=np.float64)
trans_t[k] = dense[np.ix_(loc, loc)] - c.log1p_dv
else: # every sequence empty: no states, no scatter entries, contribution 0
init_t[k] = 0.0
trans_t[k] = 0.0
return init_t, trans_t
def _markov_chain_to_value(init_hist_k: np.ndarray, trans_hist_k: np.ndarray, enc: Any, count: float) -> Any:
"""Pack the weighted histograms into the accumulator's native (init map, trans map, len value) tuple,
dropping zeros exactly as MarkovChainAccumulator.seq_update does."""
inv_key_map = enc[6]
init_map = {inv_key_map[s]: float(init_hist_k[s]) for s in range(len(inv_key_map)) if init_hist_k[s] != 0.0}
trans_map: dict[Any, dict[Any, float]] = {}
nz_a, nz_b = np.nonzero(trans_hist_k)
for a, b in zip(nz_a.tolist(), nz_b.tolist()):
trans_map.setdefault(inv_key_map[a], {})[inv_key_map[b]] = float(trans_hist_k[a, b])
return (init_map, trans_map, None) # None: the fusibility guard requires the Null length model
def _markov_chain_matches(d: Any) -> bool:
"""Match only generative chains; current closed chains therefore use the native bridge.
A chain with a Null length model is a conditional likelihood factor, not a probability law over
finite sequences. The old table path also emitted the obsolete mutable-map statistic layout.
Keeping both requirements here prevents it from being selected until a future template carries
an explicit length law and the versioned statistic schema.
"""
return (
type(d).__name__ == "MarkovChainDistribution"
and type(getattr(d, "len_dist", None)).__name__ == "NullDistribution"
and getattr(d.density_semantics(), "name", None) == "LAW"
)
register_leaf_template(
LeafTemplate(
name="markovchain",
matches=_markov_chain_matches,
data=_markov_chain_data,
params=lambda comps: {},
arity=5,
kind="chain",
chain_tables=_markov_chain_tables,
chain_to_value=_markov_chain_to_value,
)
)
def _int_categorical_table(comps: list[Any], enc: Any) -> np.ndarray:
mn = int(np.rint(np.asarray(enc).min())) if np.asarray(enc).size else 0
width = int(np.rint(np.asarray(enc).max())) - mn + 1 if np.asarray(enc).size else 1
tab = np.empty((len(comps), width), dtype=np.float64)
for k, c in enumerate(comps):
tab[k] = [float(c.log_density(mn + j)) for j in range(width)]
return tab
register_leaf_template(
LeafTemplate(
name="intcategorical",
matches=lambda d: type(d).__name__ == "IntegerCategoricalDistribution",
data=lambda enc: (_arr(enc) - int(np.rint(np.asarray(enc).min())) if np.asarray(enc).size else _arr(enc),),
params=lambda comps: {},
kind="categorical",
acc_names=("hist",),
cat_table=_int_categorical_table,
# value() = (min_val, weighted count vector over min_val .. max_val)
cat_to_value=lambda hist_k, enc, count: (
int(np.rint(np.asarray(enc).min())) if np.asarray(enc).size else 0,
hist_k.copy(),
),
)
)
def _bridge_matches(d: Any) -> bool:
"""ANY factor speaking the sequence protocol, scored and estimated through its OWN native
machinery (seq_log_density / accumulator seq_update) -- so every configuration the factor itself
supports (priors, length models, arbitrary combinators) is supported here. Registered LAST, so
every specific template outranks the bridge and this is purely the completion of coverage: a
composite with ANY protocol-bearing factor now fuses, with the un-templated factors paying
exactly their host cost while the softmax, responsibilities, and templated siblings stay in one
nopython pass."""
return (
hasattr(d, "seq_log_density")
and hasattr(d, "dist_to_encoder")
and hasattr(d, "estimator")
and callable(getattr(d, "estimator", None))
)
register_leaf_template(
LeafTemplate(
name="bridge",
matches=_bridge_matches,
data=lambda enc: (),
params=lambda comps: {},
arity=0,
kind="bridge",
)
)
# --- structure analysis ---------------------------------------------------------------------------
[docs]
@dataclass(frozen=True)
class FusedPlan:
"""A fusible model structure: K mixture components (1 if not a mixture), each a list of leaf factors."""
num_components: int
is_mixture: bool
leaf_templates: tuple[LeafTemplate, ...] # one per factor (composite factor order)
signature: tuple
component_is_composite: bool = False # whether each leaf-bearing node is a Composite (vs a bare leaf)
bare_bridge: bool = False # last-resort plan: bare combinator components bridged whole (see _node_factors)
@property
def has_matrix(self) -> bool:
"""Whether any leaf in the fused plan requires matrix BLAS handling."""
return any(t.kind == "matrix" for t in self.leaf_templates)
@property
def has_bridge(self) -> bool:
"""Whether any factor is a BRIDGED combinator (scored via its own native seq_log_density; its
sufficient statistics collected by its own accumulator, weighted by the responsibility matrix)."""
return any(t.kind == "bridge" for t in self.leaf_templates)
@property
def needs_responsibilities(self) -> bool:
"""Whether the E-step must materialize the (n, K) responsibility matrix R (matrix BLAS
accumulation and chain post-pass scatters read it after the row loop; bridged factors hand it
back to the caller for their native weighted updates)."""
return any(t.kind in ("matrix", "chain", "bridge") for t in self.leaf_templates)
def _node_factors(node: Any, bare_bridge: bool = False) -> list[Any] | None:
"""The leaf factors of a fusible *node* (an exact Composite, or a templated bare leaf), else None.
The bare-node check excludes the bridge template BY DEFAULT: a bare Mixture-of-Mixtures (no
Composite anywhere) belongs to fused_nested's in-kernel nested-tree path, which is faster than
bridging. ``bare_bridge=True`` is the LAST-RESORT pass the public entry points try only after
both the template plan and fused_nested have declined -- it lets hierarchical mixtures whose
inner components are chain/bridge-bearing composites (which neither other path can take) fuse
through per-component native scoring instead of falling all the way to the host walk."""
if type(node).__name__ == "CompositeDistribution":
return list(node.dists)
if _template_for(node, allow_bridge=bare_bridge) is not None:
return [node]
return None
[docs]
def analyze(model: Any, bare_bridge: bool = False) -> FusedPlan | None:
"""Return a :class:`FusedPlan` if ``model`` is a fusible mixture/composite of templated leaves, else None.
Handles: a single leaf, a Composite of leaves, a Mixture of leaves, and a Mixture of Composites
(every component sharing the same leaf-type structure). Anything else -> None (fall back to numpy).
The structural nodes are matched by EXACT type, not duck-typing: mixture-flavoured relatives
(SemiSupervised / Joint / Hierarchical mixtures, Select, ...) carry ``components``/``log_w`` but have
different E-step semantics, so they must keep their own kernels.
"""
tname = type(model).__name__
if tname == "MixtureDistribution":
comps = list(model.components)
per = [_node_factors(c, bare_bridge=bare_bridge) for c in comps]
if any(p is None for p in per): # a component is not a plain Composite / templated leaf
return None
templates = [_template_for(f) for f in per[0]] # type: ignore[union-attr]
if any(t is None for t in templates): # a factor has no leaf template (e.g. a Categorical)
return None
# every component must have the same leaf types in the same order
if any([_template_for(f) for f in p] != templates for p in per): # type: ignore[union-attr]
return None
comp_is_composite = type(comps[0]).__name__ == "CompositeDistribution"
sig = (
"mix",
tuple(t.fingerprint for t in templates),
comp_is_composite,
bare_bridge,
) # type: ignore[union-attr]
return FusedPlan(len(comps), True, tuple(templates), sig, comp_is_composite, bare_bridge) # type: ignore[arg-type]
# The bare-bridge allowance never applies to the ROOT itself: bridging a bare leaf/combinator
# root is one native scoring pass wrapped in a kernel -- no softmax, no templated siblings,
# nothing fused. It exists solely so a MIXTURE's components can be bridged whole (above).
factors = _node_factors(model, bare_bridge=False)
if factors is None:
return None
templates = [_template_for(f) for f in factors]
if any(t is None for t in templates): # a composite factor has no leaf template
return None
is_composite = tname == "CompositeDistribution"
sig = (
"comp",
tuple(t.fingerprint for t in templates),
is_composite,
bare_bridge,
) # type: ignore[union-attr]
return FusedPlan(1, False, tuple(templates), sig, is_composite, bare_bridge) # type: ignore[arg-type]
[docs]
def fusible(model: Any, bare_bridge: bool = True) -> bool:
"""Return whether ``model`` can use a fused scoring kernel.
``bare_bridge=False`` restricts the answer to the FAST paths (flat template, nested in-kernel):
the bare-bridge last resort scores each component through its own native ``seq_log_density``, so
it is host-speed with a fused softmax, not a single-pass win. Policy call sites that use
fusibility as a proxy for "a fused kernel would beat this schedule" (block-EM's scheduler, the
precision planner) must pass ``bare_bridge=False``; correctness call sites keep the default.
"""
if analyze(model) is not None:
return True
from mixle.stats.compute.fused_nested import fusible_nested # nested scalar trees (Mixture-of-Mixture, ...)
if fusible_nested(model):
return True
# last resort: hierarchical mixtures whose inner components carry chain/bridge factors -- the
# nested path declines them (non-scalar leaves), so bridge the inner combinators per component
return bare_bridge and analyze(model, bare_bridge=True) is not None
def _dummy(t: LeafTemplate) -> Any:
"""A throwaway leaf instance to read this template's parameter names (kept small and family-specific)."""
import mixle.stats as stats
return {
"gaussian": stats.GaussianDistribution(0.0, 1.0),
"exponential": stats.ExponentialDistribution(1.0),
"geometric": stats.GeometricDistribution(0.5),
"bernoulli": stats.BernoulliDistribution(0.5),
"poisson": stats.PoissonDistribution(1.0),
"gamma": stats.GammaDistribution(1.0, 1.0),
"loggaussian": stats.LogGaussianDistribution(0.0, 1.0),
"halfnormal": stats.HalfNormalDistribution(1.0),
"rayleigh": stats.RayleighDistribution(1.0),
"inversegaussian": stats.InverseGaussianDistribution(1.0, 1.0),
"beta": stats.BetaDistribution(2.0, 2.0),
"inversegamma": stats.InverseGammaDistribution(2.0, 1.0),
"vonmises": stats.VonMisesDistribution(0.0, 1.0),
"wrappedcauchy": stats.WrappedCauchyDistribution(0.0, 0.5),
"logseries": stats.LogSeriesDistribution(0.5),
"pareto": stats.ParetoDistribution(2.0, 1.0),
"gumbel": stats.GumbelDistribution(0.0, 1.0),
"logistic": stats.LogisticDistribution(0.0, 1.0),
"studentt": stats.StudentTDistribution(5.0, 0.0, 1.0),
"categorical": stats.CategoricalDistribution({"a": 0.5, "b": 0.5}),
"intcategorical": stats.IntegerCategoricalDistribution(0, [0.5, 0.5]),
"diaggaussian": stats.DiagonalGaussianDistribution([0.0, 0.0], [1.0, 1.0]),
"binomial": stats.BinomialDistribution(0.5, 1),
"negbinomial": stats.NegativeBinomialDistribution(1.0, 0.5),
"mvgaussian": stats.MultivariateGaussianDistribution([0.0, 0.0], [[1.0, 0.0], [0.0, 1.0]]),
"markovchain": stats.MarkovChainDistribution({"a": 1.0}, {"a": {"a": 1.0}}),
"bridge": stats.GaussianDistribution(0.0, 1.0), # params() ignores it (bridge has no static params)
}[t.name]
def _argmap(i: int, t: LeafTemplate) -> dict[str, str]:
return {pn: f"p{i}_{pn}" for pn in sorted(t.params([_dummy(t)]).keys())}
def _data_names(i: int, t: LeafTemplate) -> list[str]:
return [f"x{i}_{j}" for j in range(t.arity)]
# --- code generation ------------------------------------------------------------------------------
_COMPILED: dict[tuple, Callable] = {}
_ESTEP_COMPILED: dict[tuple, Callable] = {}
# Parallel (prange) kernel policy. Chunk count is a PURE function of n -- never of the worker count --
# and per-chunk partials combine in fixed order, so parallel results are bit-identical across runs and
# across NUMBA_NUM_THREADS (probe: fingerprints equal at 1/4/8 threads; 4.6x at 8 threads on a 2M-row
# 8-component E-step, ~1.0x at 1 thread; 3.1x end-to-end through fused_accumulate at 400k rows). Parallel
# output differs from the SEQUENTIAL kernel by float re-association: chunk-boundary reassociation in the
# E-step reductions (~1e-8 relative measured), and 1-2 ULP in the scorer -- the parallel scorer has no
# cross-row reduction, but sequential and parallel are DIFFERENT fastmath binaries, and LLVM vectorizes
# each loop nest differently (measured max 4.7e-16 relative). Bit-identity holds within each variant,
# never across variants.
_PARALLEL_MIN_OBS = (
65_536 # measured crossover ~48k on a 3-component composite (0.89x at 32k, 1.49x at 64k, 2.44x at 262k)
)
_PARALLEL_MAX_CHUNKS = 256
_PARALLEL_CHUNK_TARGET = 16_384
def _n_chunks(n: int) -> int:
return max(1, min(_PARALLEL_MAX_CHUNKS, n // _PARALLEL_CHUNK_TARGET))
def _auto_parallel(n: int) -> bool:
if n < _PARALLEL_MIN_OBS:
return False
import numba
return numba.get_num_threads() > 1
def _validated_fused_weights(weights: Any, n_rows: int) -> np.ndarray:
"""Return a finite one-dimensional weight vector aligned with encoded rows."""
try:
result = np.asarray(weights, dtype=np.float64)
except (TypeError, ValueError) as exc:
raise TypeError("fused accumulation weights must be a numeric one-dimensional array.") from exc
if result.ndim != 1:
raise ValueError("fused accumulation weights must be one-dimensional.")
if result.shape[0] != int(n_rows):
raise ValueError("fused accumulation has %d encoded rows but %d weights." % (n_rows, result.shape[0]))
if not np.all(np.isfinite(result)):
raise ValueError("fused accumulation weights must contain only finite values.")
return np.ascontiguousarray(result)
def _emit(plan: FusedPlan, acc_suffix: str = "") -> dict[str, list[str]]:
"""Per-leaf source fragments shared by the scorer and the E-step.
Each leaf contributes (by ``kind``): ``data``/``param`` argument names; ``precompute`` (matrix-only,
BLAS quad form before the row loop); ``row`` lines that add the leaf's score to ``acc`` inside the
component loop (scalar = one ``acc +=`` expr, vector = an inline per-dim loop, matrix = a precomputed
reference); E-step ``acc_args`` + ``acc`` lines (scalar/vector inline in the responsibility loop);
and matrix-only ``post`` lines (BLAS Gram accumulation after the row loop). Pure-scalar plans produce
empty ``precompute``/``post``, so their generated source is unchanged.
"""
frag: dict[str, list[str]] = {
k: [] for k in ("data_args", "param_args", "precompute", "row", "acc_args", "acc", "post")
}
for i, t in enumerate(plan.leaf_templates):
frag["data_args"].extend(_data_names(i, t))
amap = _argmap(i, t)
frag["param_args"].extend(amap.values())
if t.kind == "matrix":
frag["precompute"].extend(t.mat_precompute(i, amap)) # type: ignore[misc]
frag["row"].append("acc += " + t.mat_row(i, amap)) # type: ignore[misc]
frag["acc_args"] += [f"S1_{i}", f"S2_{i}"]
frag["post"].extend(t.mat_accumulate(i)) # type: ignore[misc]
continue
accmap = {an: f"a{i}_{an}{acc_suffix}" for an in t.acc_names}
frag["acc_args"].extend(f"a{i}_{an}" for an in t.acc_names) # signature names never carry the suffix
if t.kind == "vector":
frag["row"].extend(t.vec_row(i, amap)) # type: ignore[misc]
frag["acc"].extend(t.vec_accumulate(i, accmap)) # type: ignore[misc]
elif t.kind == "tabulated":
frag["param_args"].append(f"tab{i}") # the (K, max+1) log-pmf table, looked up by int(x)
frag["row"].append(f"acc += tab{i}[k, int(x{i}_0[i])]")
frag["acc"].append(f"{accmap['sx']}[k] += r * x{i}_0[i]")
if t.tab_hist:
frag["acc"].append(f"{accmap['hist']}[k, int(x{i}_0[i])] += r") # weighted count histogram
elif t.kind == "bridge":
# scored from a precomputed (n, K) table (each column: the factor's OWN seq_log_density for
# that component); no in-kernel statistics -- the wrapper drives the factor's native
# accumulator with the R column instead.
frag["param_args"].append(f"bt{i}")
frag["row"].append(f"acc += bt{i}[i, k]")
elif t.kind == "chain":
frag["param_args"] += [f"cinit{i}", f"ctrans{i}"] # (K,S) init and (K,S,S) transition log tables
frag["acc_args"] += [f"ih{i}", f"th{i}"] # (K,S) init and (K,S,S) transition weighted histograms
# per-row per-component chain scores, scatter-built once before the row loop:
# data arrays are x{i}_0=init row idx, x{i}_1=init state, x{i}_2=trans row idx, x{i}_3/4=from/to
frag["precompute"].extend(
[
f"cs{i} = np.zeros((n, kc))",
f"for t{i} in range(x{i}_0.shape[0]):",
" for k in range(kc):",
f" cs{i}[x{i}_0[t{i}], k] += cinit{i}[k, x{i}_1[t{i}]]",
f"for t{i} in range(x{i}_2.shape[0]):",
" for k in range(kc):",
f" cs{i}[x{i}_2[t{i}], k] += ctrans{i}[k, x{i}_3[t{i}], x{i}_4[t{i}]]",
]
)
frag["row"].append(f"acc += cs{i}[i, k]")
frag["post"].extend(
[
f"for t{i} in range(x{i}_0.shape[0]):",
" for k in range(kc):",
f" ih{i}[k, x{i}_1[t{i}]] += R[x{i}_0[t{i}], k]",
f"for t{i} in range(x{i}_2.shape[0]):",
" for k in range(kc):",
f" th{i}[k, x{i}_3[t{i}], x{i}_4[t{i}]] += R[x{i}_2[t{i}], k]",
]
)
elif t.kind == "categorical":
frag["param_args"].append(f"cat{i}") # the (K, C) log-prob table, looked up by the category index
frag["row"].append(f"acc += cat{i}[k, int(x{i}_0[i])]")
frag["acc"].append(f"{accmap['hist']}[k, int(x{i}_0[i])] += r") # weighted per-category count
else:
vals = [f"{nm}[i]" for nm in _data_names(i, t)]
frag["row"].append("acc += " + t.expr(vals, amap)) # type: ignore[misc]
frag["acc"].append(t.acc_stmt(vals, accmap, "r")) # type: ignore[misc]
return frag
import importlib.util # noqa: E402
import os # noqa: E402
import stat # noqa: E402
import sys # noqa: E402
import tempfile # noqa: E402
import threading # noqa: E402
def _default_cache_dir() -> str:
"""A per-user cache directory. A single shared ``/tmp/mixle_fused_cache`` lets another user pre-place the
module the loader imports and ``exec``s -- arbitrary code execution -- so isolate the cache by uid."""
suffix = f"_{os.getuid()}" if hasattr(os, "getuid") else ""
return os.path.join(tempfile.gettempdir(), f"mixle_fused_cache{suffix}")
_CACHE_DIR = os.environ.get("MIXLE_FUSED_CACHE_DIR") or _default_cache_dir()
_NJIT_LOCK = threading.Lock()
def _owned_privately(path: str, *, require_dir: bool) -> bool:
"""True iff ``path`` is a real dir/file we own that no other user can write (no symlink, no g/o-write).
``lstat`` (not ``stat``) so a symlink can never pass as its -- possibly attacker-owned -- target.
"""
try:
info = os.lstat(path)
except OSError:
return False
if require_dir and not stat.S_ISDIR(info.st_mode):
return False
if not require_dir and not stat.S_ISREG(info.st_mode):
return False
if hasattr(os, "getuid") and info.st_uid != os.getuid():
return False
return not (info.st_mode & (stat.S_IWGRP | stat.S_IWOTH))
def _private_cache_dir() -> str | None:
"""Return the cache dir once it is guaranteed private to us (0700, we own it), else ``None``.
``None`` tells the caller to compile in memory only -- correct result, no disk cache, no chance of
importing code from a directory another user can write.
"""
try:
os.makedirs(_CACHE_DIR, mode=0o700, exist_ok=True)
info = os.lstat(_CACHE_DIR)
# Heal a directory we own that an older (insecure) build may have created group/other-writable.
own = not hasattr(os, "getuid") or info.st_uid == os.getuid()
if stat.S_ISDIR(info.st_mode) and own and (info.st_mode & (stat.S_IWGRP | stat.S_IWOTH)):
os.chmod(_CACHE_DIR, 0o700)
except OSError:
return None
return _CACHE_DIR if _owned_privately(_CACHE_DIR, require_dir=True) else None
def _njit(src: str, fname: str, parallel: bool = False) -> Callable:
"""Compile generated source to a disk-cached ``njit`` function.
The source (``def _fused.../_estep...``) is written to a stable per-source module file and decorated
with ``numba.njit(cache=True)`` so numba persists the compiled kernel to disk: a fresh process that
fuses the same structure reuses it instead of paying the ~1-2s compile again (the matmul E-step kernel
is the costly one). ``exec``'d ``<string>`` source cannot do this -- numba's cache needs a file locator
-- so without the file the compile is paid every process. The module name is the source digest, so the
same structure maps to the same file (and cache) and different structures never collide. Any
filesystem/import problem falls back to an in-memory ``exec`` (no disk cache, identical result)."""
import numba
# The digest keys the disk-cached module. It MUST cover the decorator configuration too (the
# fastmath flag set, parallel) -- hashing the source alone silently reuses modules built under a
# previous decorator, which is exactly how the ninf/nnan fix below almost didn't ship: the old
# fastmath=True modules kept serving until the digest learned about the flags. v2 in the salt
# retires every pre-subset cache entry.
# usedforsecurity=False states what this digest is: a content key over source this module generated
# itself, used to name a cache file. It authenticates nothing -- the "never import a file we do not
# own" ownership check below is what makes loading the cached module safe, not the digest -- so a
# collision would mean a wrong-kernel cache hit, not a bypassed trust boundary. (It also keeps the
# call legal on a FIPS build, where an unqualified sha1() is refused outright.)
digest = hashlib.sha1(f"v2|parallel={parallel}|{src}".encode(), usedforsecurity=False).hexdigest()[:16]
modname = f"_pysp_fused_{digest}"
with _NJIT_LOCK:
cached_module = sys.modules.get(modname)
if cached_module is not None:
cached_function = getattr(cached_module, fname, None)
if callable(cached_function):
return cached_function
sys.modules.pop(modname, None)
path = None
try:
cache_dir = _private_cache_dir()
if cache_dir is None:
raise OSError("fused cache dir is not private to this user; compiling in memory")
path = os.path.join(cache_dir, modname + ".py")
# Write our own source unless a file we privately own is already there. Never import a
# pre-existing file we do not own: it could be attacker-planted (arbitrary code execution).
# os.replace over a symlink replaces the link itself, so this cannot be redirected outside.
if not _owned_privately(path, require_dir=False):
# fastmath SUBSET: reassociation/contraction/approximations keep the SIMD wins, but
# ninf/nnan stay OFF -- with them, LLVM folds the all--inf log-sum-exp guard
# (`m > -inf`) away and a row every component scores -inf becomes NaN instead of
# -inf (caught by the fused edge panel; the host path returns -inf).
flags = "{'reassoc', 'contract', 'arcp', 'afn', 'nsz'}"
deco = (
f"@numba.njit(fastmath={flags}, cache=True, parallel=True)"
if parallel
else f"@numba.njit(fastmath={flags}, cache=True)"
)
module_src = f"import numpy as np\nimport numba\n\n\n{deco}\n{src}\n"
tmp = f"{path}.{os.getpid()}.tmp"
with open(tmp, "w") as fh:
fh.write(module_src)
os.replace(tmp, path) # atomic publish -> concurrent processes never read a partial file
spec = importlib.util.spec_from_file_location(modname, path)
mod = importlib.util.module_from_spec(spec) # type: ignore[arg-type]
sys.modules[modname] = mod
spec.loader.exec_module(mod) # type: ignore[union-attr]
loaded = getattr(mod, fname)
if not callable(loaded):
raise TypeError("generated fused module did not define callable %s." % fname)
return loaded
except Exception: # noqa: BLE001
sys.modules.pop(modname, None)
if path is not None and _owned_privately(path, require_dir=False):
try:
os.remove(path)
except OSError:
pass
ns: dict[str, Any] = {"np": np, "numba": numba}
exec(src, ns) # noqa: S102 -- generated from fixed templates, no user input # nosec B102 # src is numba kernel source this module generated from fixed templates; this is the in-memory fallback for the identical text the disk path writes and imports, into a namespace holding only np and numba
# same fastmath subset as the disk template (ninf/nnan off: the -inf guard must hold)
return numba.njit(fastmath={"reassoc", "contract", "arcp", "afn", "nsz"}, parallel=parallel)(ns[fname])
[docs]
def clean_fused_cache(max_age_days: float = 30.0, dry_run: bool = False) -> dict[str, Any]:
"""Garbage-collect the fused kernel disk cache; returns ``{"removed": [...], "bytes": n, "dir": path}``.
The cache grows one small module (plus numba's compiled ``__pycache__`` artifacts, the bulk of the
bytes) per distinct (structure, decorator-config) forever -- nothing retires entries when a digest
salt bump or a structure change strands them. This removes cache-owned files whose modification
time is older than ``max_age_days`` (a re-used kernel is re-written/re-compiled cheaply on next
use, so an over-aggressive age costs one compile, never correctness), plus crashed writers'
``*.tmp`` leftovers after a day. Only files matching the cache's own naming pattern inside the
private cache directory are candidates; everything else is ignored. ``dry_run=True`` reports
without deleting.
"""
import glob
import time
report: dict[str, Any] = {"removed": [], "bytes": 0, "dir": None}
cache_dir = _private_cache_dir()
if cache_dir is None:
return report
report["dir"] = cache_dir
now = time.time()
cutoff = now - max_age_days * 86400.0
candidates = glob.glob(os.path.join(cache_dir, "_pysp_fused_*.py"))
candidates += glob.glob(os.path.join(cache_dir, "__pycache__", "_pysp_fused_*"))
candidates += glob.glob(os.path.join(cache_dir, "_pysp_fused_*.py.*.tmp"))
for path in candidates:
try:
info = os.lstat(path)
except OSError:
continue
if not stat.S_ISREG(info.st_mode):
continue
limit = now - 86400.0 if path.endswith(".tmp") else cutoff
if info.st_mtime >= limit:
continue
report["removed"].append(path)
report["bytes"] += int(info.st_size)
if not dry_run:
try:
os.remove(path)
except OSError:
report["removed"].pop()
report["bytes"] -= int(info.st_size)
return report
def _score_body(indent: str, row_lines: list[str], llbuf_name: str) -> list[str]:
"""The per-row score+log-sum-exp block shared by the sequential and parallel scorers."""
lines = [
f"{indent}for k in range(kc):",
f"{indent} if logw[k] == -np.inf:",
f"{indent} {llbuf_name}[k] = -np.inf",
f"{indent} else:",
f"{indent} acc = logw[k]",
]
lines += [f"{indent} {rt}" for rt in row_lines]
lines += [
f"{indent} {llbuf_name}[k] = acc",
f"{indent}invalid_evidence = False",
f"{indent}positive_infinite = 0",
f"{indent}for k in range(kc):",
f"{indent} invalid_evidence = invalid_evidence or np.isnan({llbuf_name}[k])",
f"{indent} positive_infinite += 1 if {llbuf_name}[k] == np.inf else 0",
f"{indent}m = {llbuf_name}[0]",
f"{indent}for k in range(1, kc):",
f"{indent} if {llbuf_name}[k] > m:",
f"{indent} m = {llbuf_name}[k]",
f"{indent}if invalid_evidence or positive_infinite > 1:",
f"{indent} out[i] = np.nan",
f"{indent}elif positive_infinite == 1:",
f"{indent} out[i] = np.inf",
f"{indent}elif m > -np.inf:",
f"{indent} s = 0.0",
f"{indent} for k in range(kc):",
f"{indent} s += np.exp({llbuf_name}[k] - m)",
f"{indent} out[i] = m + np.log(s)",
f"{indent}else:",
f"{indent} out[i] = -np.inf",
]
return lines
def _score_tail_quantized(indent: str, row_lines: list[str], llbuf_name: str) -> list[str]:
"""The per-row block with the log-sum-exp EXP replaced by one LUT gather over a delta-grid --
mixle.engines.qlut.quantized_logsumexp's exact semantics (round to grid, clamp the deep tail into
the bottom bin), emitted inline. Error bound: half a grid step (lse_error_bound), plus the
negligible clipped-tail term the qlut docstring quantifies."""
lines = [
f"{indent}for k in range(kc):",
f"{indent} if logw[k] == -np.inf:",
f"{indent} {llbuf_name}[k] = -np.inf",
f"{indent} else:",
f"{indent} acc = logw[k]",
]
lines += [f"{indent} {rt}" for rt in row_lines]
lines += [
f"{indent} {llbuf_name}[k] = acc",
f"{indent}invalid_evidence = False",
f"{indent}positive_infinite = 0",
f"{indent}for k in range(kc):",
f"{indent} invalid_evidence = invalid_evidence or np.isnan({llbuf_name}[k])",
f"{indent} positive_infinite += 1 if {llbuf_name}[k] == np.inf else 0",
f"{indent}m = {llbuf_name}[0]",
f"{indent}for k in range(1, kc):",
f"{indent} if {llbuf_name}[k] > m:",
f"{indent} m = {llbuf_name}[k]",
f"{indent}if invalid_evidence or positive_infinite > 1:",
f"{indent} out[i] = np.nan",
f"{indent}elif positive_infinite == 1:",
f"{indent} out[i] = np.inf",
f"{indent}elif m > -np.inf:",
f"{indent} s = 0.0",
f"{indent} for k in range(kc):",
f"{indent} qi = int(np.rint(({llbuf_name}[k] - m) * lse_inv_delta)) + lse_levm1",
f"{indent} if qi < 0:",
f"{indent} qi = 0", # scores more than `span` below the max clip into the bottom bin
f"{indent} s += lse_lut[qi]",
f"{indent} out[i] = m + np.log(s)",
f"{indent}else:",
f"{indent} out[i] = -np.inf",
]
return lines
def _compile(plan: FusedPlan, parallel: bool = False, quantized_lse: bool = False) -> Callable:
cached = _COMPILED.get((plan.signature, parallel, quantized_lse))
if cached is not None:
return cached
f = _emit(plan)
data_args = f["data_args"]
lse_args = ["lse_lut", "lse_inv_delta", "lse_levm1"] if quantized_lse else []
tail = _score_tail_quantized if quantized_lse else _score_body
if not parallel:
args = ", ".join(data_args + f["param_args"] + lse_args + ["logw", "out", "llbuf"])
lines = ["def _fused(%s):" % args, " n = out.shape[0]", " kc = logw.shape[0]"]
lines += [" " + ln for ln in f["precompute"]]
lines += [" for i in range(n):"]
lines += tail(" ", f["row"], "llbuf")
fn = _njit("\n".join(lines), "_fused")
else:
# prange over fixed chunks; out[i] rows are disjoint and each row's arithmetic is identical to
# the sequential kernel's, so the parallel scorer is BIT-IDENTICAL to it (asserted in tests).
args = ", ".join(data_args + f["param_args"] + lse_args + ["logw", "out", "n_chunks"])
lines = ["def _fused_par(%s):" % args, " n = out.shape[0]", " kc = logw.shape[0]"]
lines += [" " + ln for ln in f["precompute"]]
lines += [
" step = (n + n_chunks - 1) // n_chunks",
" for c in numba.prange(n_chunks):",
" llbuf_c = np.empty(kc)",
" for i in range(c * step, min(n, (c + 1) * step)):",
]
lines += tail(" ", f["row"], "llbuf_c")
fn = _njit("\n".join(lines), "_fused_par", parallel=True)
_COMPILED[(plan.signature, parallel, quantized_lse)] = fn
return fn
def _compile_estep(plan: FusedPlan, parallel: bool = False) -> Callable:
cached = _ESTEP_COMPILED.get((plan.signature, parallel))
if cached is not None:
return cached
f = _emit(plan, acc_suffix="_c" if parallel else "")
data_args = f["data_args"]
if not parallel:
extra = ["R"] if plan.has_bridge else []
args = ", ".join(
data_args
+ f["param_args"]
+ ["weights", "logw", "comp_counts", *f["acc_args"], *extra, "llbuf", "out_ll", "invalid_rows"]
)
lines = ["def _estep(%s):" % args, " n = invalid_rows.shape[0]", " kc = logw.shape[0]"]
if plan.needs_responsibilities and not plan.has_bridge:
lines.append(" R = np.empty((n, kc))") # read back by the matrix BLAS / chain post passes
lines += [" " + ln for ln in f["precompute"]]
lines += [
" for i in range(n):",
" wi = weights[i]",
" for k in range(kc):",
" if logw[k] == -np.inf:",
" llbuf[k] = -np.inf",
" else:",
" acc = logw[k]",
]
lines += [" " + rt for rt in f["row"]]
lines += [
" llbuf[k] = acc",
" invalid_evidence = False",
" positive_infinite = 0",
" for k in range(kc):",
" invalid_evidence = invalid_evidence or np.isnan(llbuf[k])",
" positive_infinite += 1 if llbuf[k] == np.inf else 0",
" m = llbuf[0]",
" for k in range(1, kc):",
" if llbuf[k] > m:",
" m = llbuf[k]",
" invalid_evidence = invalid_evidence or positive_infinite > 1",
" invalid_rows[i] = invalid_evidence",
" ordinary = not invalid_evidence and positive_infinite == 0 and m > -np.inf",
" s = 0.0",
" if ordinary:",
" for k in range(kc):",
" s += np.exp(llbuf[k] - m)",
# data log-likelihood, free as the posterior normalizer. A zero-weighted row must
# contribute exactly 0.0 regardless of support -- 0.0 * -inf is NaN in IEEE arithmetic, which
# would poison out_ll[0] for the whole batch on a row that's both impossible AND excluded.
" if wi != 0.0 and not invalid_evidence:",
" if positive_infinite == 1:",
" out_ll[0] += wi * np.inf",
" elif ordinary:",
" out_ll[0] += wi * (m + np.log(s))",
" else:",
" out_ll[0] += wi * -np.inf",
" for k in range(kc):",
" if invalid_evidence or (positive_infinite == 0 and not ordinary):",
" r = 0.0",
" elif positive_infinite == 1:",
" r = wi if llbuf[k] == np.inf else 0.0",
" else:",
" r = np.exp(llbuf[k] - m) / s * wi",
" comp_counts[k] += r",
]
if f["acc"]:
lines.append(" if r != 0.0:")
lines += [" " + st for st in f["acc"]]
if plan.needs_responsibilities:
lines.append(" R[i, k] = r")
lines += [" " + ln for ln in f["post"]]
fn = _njit("\n".join(lines), "_estep")
else:
# Chunk-parallel E-step. Every accumulator written inside the responsibility loop gains a leading
# n_chunks axis (allocated by the caller); chunk c only ever touches row c, so there are NO
# parallel-region reductions at all -- numba's parfor reduction analysis has nothing to reorder,
# and the caller's fixed-order sum over the chunk axis makes results bit-stable across runs and
# worker counts. Matrix leaves are untouched here: their statistics come from the sequential BLAS
# post-pass over the (row-disjoint) responsibility matrix R, exactly as in the sequential kernel.
chunked = [
f"a{i}_{an}"
for i, lt in enumerate(plan.leaf_templates)
if lt.kind not in ("matrix", "chain")
for an in lt.acc_names
]
extra = ["R"] if plan.has_bridge else []
args = ", ".join(
data_args
+ f["param_args"]
+ [
"weights",
"logw",
"comp_counts",
*f["acc_args"],
*extra,
"out_ll",
"invalid_rows",
"n_chunks",
]
)
lines = ["def _estep_par(%s):" % args, " n = invalid_rows.shape[0]", " kc = logw.shape[0]"]
if plan.needs_responsibilities and not plan.has_bridge:
lines.append(" R = np.empty((n, kc))")
lines += [" " + ln for ln in f["precompute"]]
lines += [
" step = (n + n_chunks - 1) // n_chunks",
" for c in numba.prange(n_chunks):",
" llbuf_c = np.empty(kc)",
]
lines += [f" {nm}_c = {nm}[c]" for nm in chunked]
lines += [
" for i in range(c * step, min(n, (c + 1) * step)):",
" wi = weights[i]",
" for k in range(kc):",
" if logw[k] == -np.inf:",
" llbuf_c[k] = -np.inf",
" else:",
" acc = logw[k]",
]
lines += [" " + rt for rt in f["row"]]
lines += [
" llbuf_c[k] = acc",
" invalid_evidence = False",
" positive_infinite = 0",
" for k in range(kc):",
" invalid_evidence = invalid_evidence or np.isnan(llbuf_c[k])",
" positive_infinite += 1 if llbuf_c[k] == np.inf else 0",
" m = llbuf_c[0]",
" for k in range(1, kc):",
" if llbuf_c[k] > m:",
" m = llbuf_c[k]",
" invalid_evidence = invalid_evidence or positive_infinite > 1",
" invalid_rows[i] = invalid_evidence",
" ordinary = not invalid_evidence and positive_infinite == 0 and m > -np.inf",
" s = 0.0",
" if ordinary:",
" for k in range(kc):",
" s += np.exp(llbuf_c[k] - m)",
# same zero-weight guard as the sequential kernel: 0.0 * -inf is NaN, not 0.0.
" if wi != 0.0 and not invalid_evidence:",
" if positive_infinite == 1:",
" out_ll[c] += wi * np.inf",
" elif ordinary:",
" out_ll[c] += wi * (m + np.log(s))",
" else:",
" out_ll[c] += wi * -np.inf",
" for k in range(kc):",
" if invalid_evidence or (positive_infinite == 0 and not ordinary):",
" r = 0.0",
" elif positive_infinite == 1:",
" r = wi if llbuf_c[k] == np.inf else 0.0",
" else:",
" r = np.exp(llbuf_c[k] - m) / s * wi",
" comp_counts[c, k] += r",
]
if f["acc"]:
lines.append(" if r != 0.0:")
lines += [" " + st for st in f["acc"]]
if plan.needs_responsibilities:
lines.append(" R[i, k] = r")
lines += [" " + ln for ln in f["post"]]
fn = _njit("\n".join(lines), "_estep_par", parallel=True)
_ESTEP_COMPILED[(plan.signature, parallel)] = fn
return fn
# --- marshalling ----------------------------------------------------------------------------------
def _component_factor_lists(model: Any, plan: FusedPlan) -> list[list[Any]]:
if plan.is_mixture:
return [_node_factors(c, bare_bridge=plan.bare_bridge) for c in model.components] # type: ignore[misc]
return [_node_factors(model, bare_bridge=plan.bare_bridge)] # type: ignore[list-item]
def _fused_array(value: Any, path: str) -> np.ndarray:
"""Return a contiguous numeric array safe to pass into generated code."""
try:
result = np.asarray(value)
except (TypeError, ValueError) as exc:
raise TypeError("%s must be a numeric array." % path) from exc
if result.ndim == 0:
raise ValueError("%s must have at least one dimension." % path)
if result.dtype.kind not in "biuf":
raise TypeError("%s must have a real numeric dtype." % path)
return np.ascontiguousarray(result)
def _aligned_fused_rows(current: int | None, candidate: int, path: str) -> int:
if candidate < 0:
raise ValueError("%s reports a negative row count." % path)
if current is not None and candidate != current:
raise ValueError("%s has %d rows but the fused encoding has %d." % (path, candidate, current))
return candidate if current is None else current
def _validate_fused_indices(values: np.ndarray, upper: int, path: str) -> None:
numeric = np.asarray(values)
if np.any(~np.isfinite(numeric)) or np.any(np.floor(numeric) != numeric):
raise ValueError("%s must contain finite integer indices." % path)
if numeric.size and (np.any(numeric < 0) or np.any(numeric >= upper)):
raise ValueError("%s must lie in [0, %d)." % (path, upper))
def _validated_fused_log_weights(model: Any, num_components: int) -> np.ndarray:
"""Return a one-dimensional component log-weight vector with safe values."""
result = np.asarray(getattr(model, "log_w", np.zeros(1)), dtype=np.float64)
if result.shape != (int(num_components),):
raise ValueError("fused model has %d components but log_w has shape %s." % (num_components, result.shape))
if np.any(np.isnan(result)) or np.any(np.isposinf(result)):
raise ValueError("fused model log weights may contain only finite values or -inf.")
return np.ascontiguousarray(result)
def _data_and_params(
model: Any, plan: FusedPlan, enc: Any, compute_dtype: Any = None
) -> tuple[list[np.ndarray], list[np.ndarray], dict[int, Any], int]:
"""Return (flattened data arrays, param arrays, tabulated-leaf context {i: (min_x, max_x)}).
Tabulated leaves additionally get a data-dependent ``(K, max_x+1)`` log-pmf table appended to the param
arrays (right after their -- empty -- ``params()`` block, matching the order :func:`_emit` assigns), and
their integer min/max over x recorded for ``to_value``.
When ``compute_dtype`` is a reduced float (e.g. ``np.float32``), the *floating* data and parameter
arrays are down-cast to it so the row arithmetic runs in that precision (DeepSeek-style low-precision
compute), while every accumulator -- ``acc``/``llbuf``/``out``/``comp_counts`` and the E-step
sufficient statistics -- stays float64 (high-precision accumulation), so the reduction does not drift
with N or with the number of components. Integer index arrays (categorical / tabulated lookups) are
left untouched. ``None`` keeps everything float64 (byte-identical to the legacy path).
"""
factor_lists = _component_factor_lists(model, plan)
if len(factor_lists) != plan.num_components:
raise ValueError(
"fused plan expects %d components but the model exposes %d." % (plan.num_components, len(factor_lists))
)
for component, factors in enumerate(factor_lists):
if factors is None or len(factors) != len(plan.leaf_templates):
count = 0 if factors is None else len(factors)
raise ValueError(
"fused component %d exposes %d factors but the plan expects %d."
% (component, count, len(plan.leaf_templates))
)
# A Composite encodes as a per-factor tuple; a bare leaf encodes as its own payload (which may itself
# be a tuple, e.g. Poisson's (x, lgamma(x+1))) -- so key on the structure, not isinstance(enc, tuple).
if plan.component_is_composite:
if not isinstance(enc, (tuple, list)):
raise TypeError("fused composite encoding must be a sequence with one payload per factor.")
factor_encs = tuple(enc)
else:
factor_encs = (enc,)
if len(factor_encs) != len(plan.leaf_templates):
raise ValueError(
"fused encoding has %d factor payloads but the plan expects %d."
% (len(factor_encs), len(plan.leaf_templates))
)
data_arrays: list[np.ndarray] = []
param_arrays: list[np.ndarray] = []
tab_ctx: dict[int, Any] = {}
n_rows: int | None = None # chain data arrays are transition-length, so rows come from another source
for i, t in enumerate(plan.leaf_templates):
try:
raw_arrays = t.data(factor_encs[i])
except (IndexError, TypeError, ValueError) as exc:
raise ValueError("fused factor %d (%s) has an invalid encoded payload: %s" % (i, t.name, exc)) from exc
if not isinstance(raw_arrays, (tuple, list)) or len(raw_arrays) != t.arity:
count = len(raw_arrays) if isinstance(raw_arrays, (tuple, list)) else -1
raise ValueError(
"fused factor %d (%s) produced %d data arrays but its arity is %d." % (i, t.name, count, t.arity)
)
arrs = tuple(
_fused_array(value, "fused factor %d (%s) data[%d]" % (i, t.name, j)) for j, value in enumerate(raw_arrays)
)
if t.kind in ("scalar", "tabulated", "categorical") and any(array.ndim != 1 for array in arrs):
raise ValueError("fused factor %d (%s) requires one-dimensional data arrays." % (i, t.name))
if t.kind in ("vector", "matrix") and (len(arrs) != 1 or arrs[0].ndim != 2):
raise ValueError("fused factor %d (%s) requires one two-dimensional data array." % (i, t.name))
if t.kind == "chain" and any(array.ndim != 1 for array in arrs):
raise ValueError("fused factor %d (%s) requires one-dimensional scatter arrays." % (i, t.name))
data_arrays.extend(arrs) # arity arrays, flattened in leaf order
if t.kind == "chain":
chain_enc = factor_encs[i]
if not isinstance(chain_enc, (tuple, list)) or len(chain_enc) != 8:
raise ValueError("fused chain encoding must contain exactly eight fields.")
chain_count = chain_enc[0]
if (
isinstance(chain_count, (bool, np.bool_))
or not np.isscalar(chain_count)
or not np.isfinite(chain_count)
or int(chain_count) != chain_count
or int(chain_count) < 0
):
raise ValueError("fused chain row count must be a non-negative integer.")
n_rows = _aligned_fused_rows(n_rows, int(chain_count), "fused factor %d (%s)" % (i, t.name))
if len(arrs[0]) != len(arrs[1]):
raise ValueError("fused chain initial row and state arrays must have equal lengths.")
if not (len(arrs[2]) == len(arrs[3]) == len(arrs[4])):
raise ValueError("fused chain transition row/from/to arrays must have equal lengths.")
_validate_fused_indices(arrs[0], int(chain_count), "fused chain initial row indices")
_validate_fused_indices(arrs[2], int(chain_count), "fused chain transition row indices")
state_count = len(chain_enc[6])
_validate_fused_indices(arrs[1], state_count, "fused chain initial state indices")
_validate_fused_indices(arrs[3], state_count, "fused chain previous-state indices")
_validate_fused_indices(arrs[4], state_count, "fused chain next-state indices")
else:
for j, array in enumerate(arrs):
n_rows = _aligned_fused_rows(
n_rows,
int(array.shape[0]),
"fused factor %d (%s) data[%d]" % (i, t.name, j),
)
comps_i = [factor_lists[k][i] for k in range(plan.num_components)] # type: ignore[index]
try:
pdict = t.params(comps_i)
except (IndexError, TypeError, ValueError) as exc:
raise ValueError(
"fused factor %d (%s) has incompatible component parameter geometry: %s" % (i, t.name, exc)
) from exc
if not isinstance(pdict, dict) or any(not isinstance(name, str) or not name.isidentifier() for name in pdict):
raise TypeError("fused factor %d (%s) params must be a mapping with identifier keys." % (i, t.name))
factor_params = []
for name in sorted(pdict):
parameter = _fused_array(pdict[name], "fused factor %d (%s) parameter %s" % (i, t.name, name))
if parameter.shape[0] != plan.num_components:
raise ValueError(
"fused factor %d (%s) parameter %s has component width %d; expected %d."
% (i, t.name, name, parameter.shape[0], plan.num_components)
)
factor_params.append(parameter)
if t.kind in ("vector", "matrix"):
dimension = int(arrs[0].shape[1])
for parameter in factor_params:
if parameter.ndim >= 2 and parameter.shape[1] != dimension:
raise ValueError(
"fused factor %d (%s) parameter shape %s is incompatible with data dimension %d."
% (i, t.name, parameter.shape, dimension)
)
if t.kind == "matrix" and parameter.ndim == 3 and parameter.shape[2] != dimension:
raise ValueError(
"fused factor %d (%s) matrix parameter shape %s is incompatible with data dimension %d."
% (i, t.name, parameter.shape, dimension)
)
param_arrays.extend(factor_params)
if t.kind == "bridge":
cols = []
for component, distribution in enumerate(comps_i):
column = _fused_array(
distribution.seq_log_density(factor_encs[i]),
"fused bridge factor %d component %d scores" % (i, component),
)
if column.ndim != 1:
raise ValueError("fused bridge score columns must be one-dimensional.")
n_rows = _aligned_fused_rows(
n_rows,
int(column.shape[0]),
"fused bridge factor %d component %d" % (i, component),
)
cols.append(np.asarray(column, dtype=np.float64))
table = np.ascontiguousarray(np.stack(cols, axis=1)) # (n, K)
if table.shape != (int(n_rows), plan.num_components):
raise ValueError(
"fused bridge table has shape %s; expected (%d, %d)." % (table.shape, n_rows, plan.num_components)
)
param_arrays.append(table)
tab_ctx[i] = (factor_encs[i], 0) # the factor's own encoding, for the wrapper's native updates
if t.kind == "chain":
init_table, trans_table = t.chain_tables(comps_i, factor_encs[i]) # type: ignore[misc]
init_table = _fused_array(init_table, "fused chain initial table")
trans_table = _fused_array(trans_table, "fused chain transition table")
state_width = max(len(factor_encs[i][6]), 1)
if init_table.shape != (plan.num_components, state_width):
raise ValueError(
"fused chain initial table has shape %s; expected (%d, %d)."
% (init_table.shape, plan.num_components, state_width)
)
if trans_table.shape != (plan.num_components, state_width, state_width):
raise ValueError(
"fused chain transition table has shape %s; expected (%d, %d, %d)."
% (trans_table.shape, plan.num_components, state_width, state_width)
)
param_arrays.extend((init_table, trans_table))
tab_ctx[i] = (factor_encs[i], state_width) # (encoding for to_value, S states)
if t.kind == "tabulated":
x = arrs[0]
if np.any(~np.isfinite(x)) or np.any(np.floor(x) != x) or np.any(x < 0):
raise ValueError(
"fused tabulated factor %d (%s) requires finite non-negative integer data." % (i, t.name)
)
mx = int(np.rint(x.max())) if x.size else 0
mn = int(np.rint(x.min())) if x.size else 0
table = _fused_array(t.tab_table(comps_i, mx), "fused factor %d (%s) table" % (i, t.name)) # type: ignore[misc]
if table.shape != (plan.num_components, mx + 1):
raise ValueError(
"fused factor %d (%s) table has shape %s; expected (%d, %d)."
% (i, t.name, table.shape, plan.num_components, mx + 1)
)
param_arrays.append(table)
tab_ctx[i] = (mn, mx)
elif t.wants_minmax:
x = arrs[0]
tab_ctx[i] = (float(x.min()) if x.size else 0.0, float(x.max()) if x.size else 0.0)
elif t.kind == "categorical":
table = _fused_array(
t.cat_table(comps_i, factor_encs[i]), # type: ignore[misc]
"fused factor %d (%s) category table" % (i, t.name),
)
if table.ndim != 2 or table.shape[0] != plan.num_components:
raise ValueError(
"fused factor %d (%s) category table must have shape (%d, C), got %s."
% (i, t.name, plan.num_components, table.shape)
)
_validate_fused_indices(
arrs[0],
table.shape[1],
"fused factor %d (%s) category indices" % (i, t.name),
)
param_arrays.append(table)
tab_ctx[i] = (factor_encs[i], table.shape[1]) # (encoding for to_value, C for histogram width)
if n_rows is None:
raise ValueError("fused encoding does not expose an observation row count.")
if compute_dtype is not None and np.dtype(compute_dtype) != np.float64:
if np.dtype(compute_dtype) != np.float32:
raise ValueError(
"fused kernels support reduced precision only in float32: numba cannot compile float16 / "
"bfloat16 or sub-byte formats on CPU. Got %r." % (compute_dtype,)
)
# Down-cast only the FLOATING arrays; integer index/lookup arrays must keep their dtype. The
# generated kernels promote back to the float64 accumulators on every ``acc += ...``.
cast = lambda a: np.ascontiguousarray(a, dtype=compute_dtype) if a.dtype.kind == "f" else a # noqa: E731
data_arrays = [cast(a) for a in data_arrays]
param_arrays = [cast(a) for a in param_arrays]
return data_arrays, param_arrays, tab_ctx, int(n_rows if n_rows is not None else 0)
[docs]
def fused_seq_log_density(
model: Any,
enc: Any,
compute_dtype: Any = None,
parallel: bool | None = None,
lse_bits: int | None = None,
lse_span: float = 24.0,
) -> np.ndarray:
"""Per-row log densities of ``model`` over encoding ``enc`` via one fused numba pass.
``compute_dtype`` (e.g. ``np.float32``) runs the row arithmetic in reduced precision while the
log-sum-exp accumulator and output stay float64; ``None`` keeps the byte-identical float64 path.
``parallel``: ``None`` auto-engages the chunked prange scorer for large inputs (>= 65536 rows -- the measured crossover -- and
more than one numba worker); ``True``/``False`` force it. Rows are disjoint (no cross-row
reduction), so the parallel scorer is bit-identical across runs and worker counts; against the
sequential kernel it agrees to 1-2 ULP (different fastmath binaries vectorize differently --
measured max 4.7e-16 relative), not bit-for-bit.
``lse_bits`` (OPT-IN, banded): replace the per-row log-sum-exp's exp calls with one gather from a
``2**lse_bits``-entry table over a ``lse_span/2**lse_bits`` grid -- the qlut quantized-LSE kernel
inlined. Per-row error is bounded by ``mixle.engines.qlut.lse_error_bound(lse_bits, lse_span)``
(12 bits ~ 2.9e-3 bound, ~6e-5 measured), so this is for scoring/audit paths whose compute band
tolerates a delta-grid; E-steps deliberately keep exact exp (a perturbed objective could flip
monotone-gate accepts). ``None`` (default) is the exact path, byte-identical to before.
Raises ``ValueError`` if ``model`` is not fusible -- callers should check :func:`fusible` first.
"""
plan = analyze(model)
if plan is None:
from mixle.stats.compute.fused_nested import fused_nested_seq_log_density, fusible_nested
if fusible_nested(model):
# nested scalar tree; same dtype/parallel/quantized-LSE contract (nested error bound
# compounds per mixture level -- see the nested wrapper's docstring)
return fused_nested_seq_log_density(
model, enc, compute_dtype=compute_dtype, parallel=parallel, lse_bits=lse_bits, lse_span=lse_span
)
plan = analyze(model, bare_bridge=True) # hierarchical mixtures over chain/bridge composites
if plan is None:
raise ValueError("%s is not fusible on any path (template, nested, bridge)." % type(model).__name__)
data_arrays, param_arrays, _, n = _data_and_params(model, plan, enc, compute_dtype)
logw = _validated_fused_log_weights(model, plan.num_components)
out = np.empty(n, dtype=np.float64)
if parallel is None:
parallel = _auto_parallel(n)
lse_extra: tuple = ()
quantized = lse_bits is not None
if quantized:
if not 1 <= int(lse_bits) <= 24:
raise ValueError(f"need 1 <= lse_bits <= 24, got {lse_bits}")
if lse_span <= 0:
raise ValueError(f"lse_span must be positive, got {lse_span}")
levels = 1 << int(lse_bits)
delta = float(lse_span) / levels
table = np.exp((np.arange(levels, dtype=np.float64) - (levels - 1)) * delta)
lse_extra = (table, 1.0 / delta, levels - 1)
if parallel:
_compile(plan, parallel=True, quantized_lse=quantized)(
*data_arrays, *param_arrays, *lse_extra, logw, out, _n_chunks(n)
)
else:
llbuf = np.empty(plan.num_components, dtype=np.float64)
_compile(plan, quantized_lse=quantized)(*data_arrays, *param_arrays, *lse_extra, logw, out, llbuf)
raise_for_invalid_log_evidence(out)
return out
# --- fused E-step (score + responsibilities + per-leaf weighted sufficient statistics, one njit) ----
[docs]
def fusible_estep(model: Any, bare_bridge: bool = True) -> bool:
"""Return whether ``model`` can use a fused E-step accumulation kernel.
``bare_bridge=False`` restricts to the fast paths, exactly as in :func:`fusible`.
"""
plan = analyze(model)
if plan is None:
from mixle.stats.compute.fused_nested import fusible_nested # nested scalar trees fit the E-step too
if fusible_nested(model):
return True
plan = analyze(model, bare_bridge=True) if bare_bridge else None
if plan is None:
return False
hook = {
"scalar": lambda t: t.acc_stmt,
"vector": lambda t: t.vec_accumulate,
"matrix": lambda t: t.mat_accumulate,
"tabulated": lambda t: t.tab_table,
"categorical": lambda t: t.cat_table,
"chain": lambda t: t.chain_to_value,
"bridge": lambda t: t.matches, # support == matching: statistics come from the factor's own accumulator
}
return all(hook[t.kind](t) is not None for t in plan.leaf_templates)
[docs]
def fused_accumulate(
model: Any,
enc: Any,
weights: np.ndarray,
return_ll: bool = False,
compute_dtype: Any = None,
parallel: bool | None = None,
) -> Any:
"""Run one fused E-step and return the sufficient statistic in the estimator's ``value()`` format.
The whole E-step -- component scoring, responsibility soft-max, and per-leaf weighted-statistic
accumulation (scalar inline, matrix via BLAS) -- runs in a single nopython pass, then is packed into
the exact tuple shape ``estimate(nobs, suff_stat)`` expects. With ``return_ll`` the weighted data
log-likelihood (the posterior normalizer, computed for free in the same pass) is also returned as
``(suff_stat, ll)`` so the EM loop can skip a separate scoring pass. Raises ``ValueError`` if not
fusible.
``parallel``: ``None`` auto-engages the chunked prange E-step for large inputs (>= 65536 rows -- the measured crossover -- and
more than one numba worker); ``True``/``False`` force it. Chunk boundaries are a pure function of n
and chunk partials combine in fixed order, so parallel results are bit-identical across runs AND
across worker counts; they differ from the sequential kernel only by float re-association across
chunk boundaries (~1e-7 relative). Matrix leaves accumulate through the same sequential BLAS
post-pass in both variants.
"""
plan = analyze(model)
if plan is None:
from mixle.stats.compute.fused_nested import fused_nested_accumulate, fusible_nested
if fusible_nested(model):
return fused_nested_accumulate(
model, enc, weights, return_ll=return_ll, compute_dtype=compute_dtype, parallel=parallel
) # nested scalar tree
plan = analyze(model, bare_bridge=True) # hierarchical mixtures over chain/bridge composites
if plan is None:
raise ValueError("%s is not fusible on any path (template, nested, bridge)." % type(model).__name__)
if not fusible_estep(model):
raise ValueError("%s is not a fusible E-step (an unsupported leaf)." % type(model).__name__)
K = plan.num_components
data_arrays, param_arrays, tab_ctx, n = _data_and_params(model, plan, enc, compute_dtype)
weights = _validated_fused_weights(weights, n)
if parallel is None:
parallel = _auto_parallel(n)
nc = _n_chunks(n) if parallel else 0
chunk = (lambda shape: (nc, *shape)) if parallel else (lambda shape: shape)
# per-leaf accumulator arrays, in the same leaf order the generated signature expects. data_arrays is
# flattened by arity, so track the running offset to find each leaf's first data array.
scalar_acc: list[dict[str, np.ndarray]] = []
matrix_acc: list[tuple[np.ndarray, np.ndarray]] = []
chain_acc: list[tuple[np.ndarray, np.ndarray]] = []
acc_arrays: list[np.ndarray] = []
offset = 0
for i, t in enumerate(plan.leaf_templates):
if t.kind == "bridge": # no in-kernel statistics; arity 0
scalar_acc.append({})
matrix_acc.append((np.empty(0), np.empty(0)))
chain_acc.append((np.empty(0), np.empty(0)))
continue
if t.kind == "chain":
S = tab_ctx[i][1]
# written by the sequential post-pass over R in BOTH kernel variants, so never chunked
ih = np.zeros((K, S), dtype=np.float64)
th = np.zeros((K, S, S), dtype=np.float64)
chain_acc.append((ih, th))
scalar_acc.append({})
matrix_acc.append((np.empty(0), np.empty(0)))
acc_arrays += [ih, th]
offset += t.arity
continue
chain_acc.append((np.empty(0), np.empty(0)))
if t.kind == "matrix":
d = data_arrays[offset].shape[1]
s1 = np.zeros((K, d), dtype=np.float64)
s2 = np.zeros((K, d, d), dtype=np.float64)
matrix_acc.append((s1, s2))
scalar_acc.append({})
acc_arrays += [s1, s2]
elif t.kind == "vector":
d = data_arrays[offset].shape[1] # (K,D) per-dim weighted statistics
ad = {an: np.zeros(chunk((K, d)), dtype=np.float64) for an in t.acc_names}
scalar_acc.append(ad)
matrix_acc.append((np.empty(0), np.empty(0)))
acc_arrays.extend(ad.values())
elif t.kind == "tabulated":
width = tab_ctx[i][1] + 1 # max_x + 1; sx is (K,), the histogram (if any) is (K, max_x+1)
ad = {
an: np.zeros(chunk((K, width)) if an == "hist" else chunk((K,)), dtype=np.float64) for an in t.acc_names
}
scalar_acc.append(ad)
matrix_acc.append((np.empty(0), np.empty(0)))
acc_arrays.extend(ad.values())
elif t.kind == "categorical":
width = tab_ctx[i][1] # C categories; the only statistic is the (K, C) weighted count histogram
ad = {"hist": np.zeros(chunk((K, width)), dtype=np.float64)}
scalar_acc.append(ad)
matrix_acc.append((np.empty(0), np.empty(0)))
acc_arrays.append(ad["hist"])
else:
# min-combined accumulators (per-component support minima) start at +inf, additive ones at 0
ad = {
an: np.full(chunk((K,)), np.inf, dtype=np.float64)
if an in t.min_acc_names
else np.zeros(chunk((K,)), dtype=np.float64)
for an in t.acc_names
}
scalar_acc.append(ad)
matrix_acc.append((np.empty(0), np.empty(0)))
acc_arrays.extend(ad.values())
offset += t.arity
logw = _validated_fused_log_weights(model, K)
bridge_R = [np.empty((n, K), dtype=np.float64)] if plan.has_bridge else []
if parallel:
comp_counts = np.zeros((nc, K), dtype=np.float64)
out_ll = np.zeros(nc, dtype=np.float64)
invalid_rows = np.zeros(n, dtype=np.bool_)
_compile_estep(plan, parallel=True)(
*data_arrays,
*param_arrays,
weights,
logw,
comp_counts,
*acc_arrays,
*bridge_R,
out_ll,
invalid_rows,
nc,
)
# fixed-order combine over the chunk axis (numpy's single-threaded pairwise reduce): the SAME
# partials in the SAME order regardless of how many workers computed them -> bit-stable results.
comp_counts = comp_counts.sum(axis=0)
for i, t in enumerate(plan.leaf_templates):
if t.kind != "matrix":
for an in list(scalar_acc[i]):
# support minima are non-additive: chunk partials must min-combine, never sum
chunk_op = np.minimum.reduce if an in t.min_acc_names else np.add.reduce
scalar_acc[i][an] = chunk_op(scalar_acc[i][an], axis=0)
else:
comp_counts = np.zeros(K, dtype=np.float64)
llbuf = np.empty(K, dtype=np.float64)
out_ll = np.zeros(1, dtype=np.float64)
invalid_rows = np.zeros(n, dtype=np.bool_)
_compile_estep(plan)(
*data_arrays,
*param_arrays,
weights,
logw,
comp_counts,
*acc_arrays,
*bridge_R,
llbuf,
out_ll,
invalid_rows,
)
if invalid_rows.any():
raise InvalidMixtureEvidenceError(
np.flatnonzero(invalid_rows),
"component scores contain NaN or ambiguous +inf evidence",
)
bridge_values: dict[int, list[Any]] = {}
if plan.has_bridge:
# Each bridged factor's sufficient statistics come from its OWN accumulator, weighted by the
# responsibility column the kernel just computed -- byte-for-byte the host mixture E-step's
# semantics for that factor (priors and all), with everything else still fused.
factor_lists = _component_factor_lists(model, plan)
for i, t in enumerate(plan.leaf_templates):
if t.kind != "bridge":
continue
enc_i = tab_ctx[i][0]
vals = []
for k in range(K):
factor_k = factor_lists[k][i]
acc = factor_k.estimator().accumulator_factory().make()
acc.seq_update(enc_i, np.ascontiguousarray(bridge_R[0][:, k]), factor_k)
vals.append(acc.value())
bridge_values[i] = vals
def leaf_value(i: int, t: LeafTemplate, k: int) -> Any:
if t.kind == "bridge":
return bridge_values[i][k]
if t.kind == "chain":
ih, th = chain_acc[i]
enc_i, _ = tab_ctx[i]
return t.chain_to_value(ih[k], th[k], enc_i, float(comp_counts[k])) # type: ignore[misc]
if t.kind == "matrix":
s1, s2 = matrix_acc[i]
return t.to_value((s1[k], s2[k]), float(comp_counts[k])) # type: ignore[misc]
if t.kind == "tabulated":
mn, mx = tab_ctx[i]
hist_k = scalar_acc[i]["hist"][k] if t.tab_hist else None
return t.tab_to_value(float(scalar_acc[i]["sx"][k]), float(comp_counts[k]), hist_k, mn, mx) # type: ignore[misc]
if t.kind == "categorical":
enc_i, _ = tab_ctx[i]
return t.cat_to_value(scalar_acc[i]["hist"][k], enc_i, float(comp_counts[k])) # type: ignore[misc]
stats_k = tuple(scalar_acc[i][an][k] for an in t.acc_names)
if t.wants_minmax:
mn, mx = tab_ctx[i]
return t.to_value_g(stats_k, float(comp_counts[k]), mn, mx) # type: ignore[misc]
return t.to_value(stats_k, float(comp_counts[k])) # type: ignore[misc]
def node_value(k: int) -> Any:
leaf_vals = [leaf_value(i, t, k) for i, t in enumerate(plan.leaf_templates)]
return tuple(leaf_vals) if plan.component_is_composite else leaf_vals[0]
suff = (comp_counts, tuple(node_value(k) for k in range(K))) if plan.is_mixture else node_value(0)
return (suff, float(out_ll.sum())) if return_ll else suff
# --- kernel wiring (used by optimize(..., engine=FUSED_NUMPY_ENGINE) for fusible models) ------------
[docs]
class FusedKernel:
"""A duck-typed ``Kernel`` backed by the source-generated fused scorer and E-step."""
def __init__(self, dist: Any, engine: Any, estimator: Any = None) -> None:
self.dist = dist
self.engine = engine
self.estimator = estimator
# Reduced-precision row arithmetic when the engine declares a non-float64 float dtype (e.g. a
# ``NumpyEngine(dtype="float32", prefer_fused=True)``); None keeps the byte-identical float64 path.
# The default fused engine carries dtype=None, so auto-fusion never silently lowers precision.
edt = getattr(engine, "dtype", None)
self.compute_dtype = edt if edt is not None and np.dtype(edt) != np.float64 else None
self.last_ll: float | None = None # data LL of the last accumulate() pass (posterior normalizer)
# Tuning knobs (``optimize(fused_options=...)`` sets these). ``parallel`` overrides the
# observation-count auto-gate in both directions; ``lse_bits``/``lse_span`` opt scoring into
# quantized LSE. E-steps stay exact by design, so accumulate() never sees the lse knobs.
self.parallel: bool | None = None
self.lse_bits: int | None = None
self.lse_span: float = 24.0
[docs]
def encode(self, data: Any) -> Any:
"""Encode raw data through the distribution's sequence encoder."""
return self.dist.dist_to_encoder().seq_encode(data)
[docs]
def score(self, enc: Any) -> np.ndarray:
"""Score encoded data with the fused log-density kernel."""
return fused_seq_log_density(
self.dist,
getattr(enc, "engine_payload", enc),
self.compute_dtype,
parallel=self.parallel,
lse_bits=self.lse_bits,
lse_span=self.lse_span,
)
[docs]
def accumulate(self, enc: Any, weights: Any) -> Any:
"""Accumulate weighted sufficient statistics with the fused E-step kernel."""
w = np.asarray(self.engine.to_numpy(weights) if hasattr(self.engine, "to_numpy") else weights, dtype=np.float64)
suff, self.last_ll = fused_accumulate(
self.dist,
getattr(enc, "engine_payload", enc),
w,
return_ll=True,
compute_dtype=self.compute_dtype,
parallel=self.parallel,
)
return suff
[docs]
def refresh(self, dist: Any) -> None:
"""Refresh the kernel wrapper with a newly estimated distribution."""
self.dist = dist
[docs]
class FusedKernelFactory:
"""Build a :class:`FusedKernel` on a numba-capable engine for fusible models, else delegate.
Scoring needs only :func:`fusible`; estimation also needs the fused E-step (:func:`fusible_estep`),
so when an estimator is present we require the stronger check. Everything else falls through to the
declaration/numba/generic factory (the previous behaviour), so registering this is never a regression.
"""
def __init__(self, fallback: Any = None) -> None:
self._fallback = fallback
def _fallback_factory(self) -> Any:
if self._fallback is None:
from mixle.stats.compute.kernel import GeneratedNumbaKernelFactory
self._fallback = GeneratedNumbaKernelFactory()
return self._fallback
[docs]
def build(self, dist: Any, engine: Any, estimator: Any = None) -> Any:
"""Build a fused kernel when supported, otherwise delegate to the fallback factory."""
ok = fusible_estep(dist) if estimator is not None else fusible(dist)
if getattr(engine, "supports_numba", False) and ok:
return FusedKernel(dist, engine, estimator=estimator)
return self._fallback_factory().build(dist, engine, estimator=estimator)
__all__ = [
"LeafTemplate",
"register_leaf_template",
"analyze",
"fusible",
"fusible_estep",
"fused_seq_log_density",
"fused_accumulate",
"clean_fused_cache",
"FusedKernel",
"FusedKernelFactory",
"FusedPlan",
]