mixle.stats.compute.fused_codegen module

Source-generated fused numba kernels for composite / mixture trees.

The generated-numba kernels in 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.

Generated kernels are compiled once and disk-cached (see _njit()), so the compile cost is paid once per structure ever, not per process. A leaf with no template (e.g. Laplace, whose weighted-median MLE keeps the raw observations) -> fusible() is False -> numpy.

class LeafTemplate(name, matches, data, params, to_value=None, arity=1, kind='scalar', expr=None, acc_names=(), acc_stmt=None, mat_precompute=None, mat_row=None, mat_accumulate=None, vec_row=None, vec_accumulate=None, tab_table=None, tab_hist=False, tab_to_value=None, wants_minmax=False, to_value_g=None, cat_table=None, cat_to_value=None, dtype='float64')

Bases: object

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.

Parameters:

• name (str)

• matches (Callable[[Any], bool])

• data (Callable[[Any], tuple[ndarray, ...]])

• params (Callable[[list[Any]], dict[str, ndarray]])

• to_value (Callable[[tuple, float], tuple] | None)

• arity (int)

• kind (str)

• expr (Callable[[list[str], dict[str, str]], str] | None)

• acc_names (tuple[str, ...])

• acc_stmt (Callable[[list[str], dict[str, str], str], str] | None)

• mat_precompute (Callable[[int, dict[str, str]], list[str]] | None)

• mat_row (Callable[[int, dict[str, str]], str] | None)

• mat_accumulate (Callable[[int], list[str]] | None)

• vec_row (Callable[[int, dict[str, str]], list[str]] | None)

• vec_accumulate (Callable[[int, dict[str, str]], list[str]] | None)

• tab_table (Callable[[list[Any], int], ndarray] | None)

• tab_hist (bool)

• tab_to_value (Callable[[float, float, ndarray | None, int, int], tuple] | None)

• wants_minmax (bool)

• to_value_g (Callable[[tuple, float, float, float], tuple] | None)

• cat_table (Callable[[list[Any], Any], ndarray] | None)

• cat_to_value (Callable[[ndarray, Any, float], Any] | None)

• dtype (str)

register_leaf_template(t)

Register a leaf template for fused scoring and E-step generation.

Parameters:

t (LeafTemplate)

Return type:

None

analyze(model)

Return a 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.

Parameters:

model (Any)

Return type:

FusedPlan | None

fusible(model)

Return whether model can use a fused scoring kernel.

Parameters:

model (Any)

Return type:

bool

fusible_estep(model)

Return whether model can use a fused E-step accumulation kernel.

Parameters:

model (Any)

Return type:

bool

fused_seq_log_density(model, enc, compute_dtype=None)

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.

Raises ValueError if model is not fusible – callers should check fusible() first.

Parameters:

• model (Any)

• enc (Any)

• compute_dtype (Any)

Return type:

ndarray

fused_accumulate(model, enc, weights, return_ll=False, compute_dtype=None)

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.

Parameters:

• model (Any)

• enc (Any)

• weights (ndarray)

• return_ll (bool)

• compute_dtype (Any)

Return type:

Any

class FusedKernel(dist, engine, estimator=None)

Bases: object

A duck-typed Kernel backed by the source-generated fused scorer and E-step.

Parameters:

• dist (Any)

• engine (Any)

• estimator (Any)

encode(data)

Encode raw data through the distribution’s sequence encoder.

Parameters:

data (Any)

Return type:

Any

score(enc)

Score encoded data with the fused log-density kernel.

Parameters:

enc (Any)

Return type:

ndarray

accumulate(enc, weights)

Accumulate weighted sufficient statistics with the fused E-step kernel.

Parameters:

• enc (Any)

• weights (Any)

Return type:

Any

refresh(dist)

Refresh the kernel wrapper with a newly estimated distribution.

Parameters:

dist (Any)

Return type:

None

class FusedKernelFactory(fallback=None)

Bases: object

Build a FusedKernel on a numba-capable engine for fusible models, else delegate.

Scoring needs only fusible(); estimation also needs the fused E-step (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.

Parameters:

fallback (Any)

build(dist, engine, estimator=None)

Build a fused kernel when supported, otherwise delegate to the fallback factory.

Parameters:

• dist (Any)

• engine (Any)

• estimator (Any)

Return type:

Any

class FusedPlan(num_components, is_mixture, leaf_templates, signature, component_is_composite=False)

Bases: object

A fusible model structure: K mixture components (1 if not a mixture), each a list of leaf factors.

Parameters:

• num_components (int)

• is_mixture (bool)

• leaf_templates (tuple[LeafTemplate, ...])

• signature (tuple)

• component_is_composite (bool)

property has_matrix: bool

Whether any leaf in the fused plan requires matrix BLAS handling.