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.

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 (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 _njit()), so the compile cost is paid once per structure ever, not per process; clean_fused_cache() garbage-collects entries stranded by structure changes or digest-salt bumps.

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, min_acc_names=(), cat_table=None, cat_to_value=None, chain_tables=None, chain_to_value=None, dtype='float64', cache_version='1')

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)

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

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

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

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

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

• dtype (str)

• cache_version (str)

register_leaf_template(t)

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).

Parameters:

t (LeafTemplate)

Return type:

None

analyze(model, bare_bridge=False)

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)

• bare_bridge (bool)

Return type:

FusedPlan | None

fusible(model, bare_bridge=True)

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.

Parameters:

• model (Any)

• bare_bridge (bool)

Return type:

bool

fusible_estep(model, bare_bridge=True)

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

bare_bridge=False restricts to the fast paths, exactly as in fusible().

Parameters:

• model (Any)

• bare_bridge (bool)

Return type:

bool

fused_seq_log_density(model, enc, compute_dtype=None, parallel=None, lse_bits=None, lse_span=24.0)

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 fusible() first.

Parameters:

• model (Any)

• enc (Any)

• compute_dtype (Any)

• parallel (bool | None)

• lse_bits (int | None)

• lse_span (float)

Return type:

ndarray

fused_accumulate(model, enc, weights, return_ll=False, compute_dtype=None, parallel=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.

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.

Parameters:

• model (Any)

• enc (Any)

• weights (ndarray)

• return_ll (bool)

• compute_dtype (Any)

• parallel (bool | None)

Return type:

Any

clean_fused_cache(max_age_days=30.0, dry_run=False)

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.

Parameters:

• max_age_days (float)

• dry_run (bool)

Return type:

dict[str, 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, bare_bridge=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)

• bare_bridge (bool)

property has_matrix: bool

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

property has_bridge: 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).

property needs_responsibilities: 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).