mixle.stats.compute.fused_nested module¶
Recursive fusion for arbitrarily nested Composite / Mixture trees of scalar leaves.
The flat fused_codegen handles depth-2 (Mixture -> Composite -> leaf) with a component loop. This
module handles arbitrary nesting – a Composite factor that is itself a Mixture, a Mixture of Mixtures, a
Mixture of Composites whose factors nest, … – by UNROLLING the static tree into straight-line numba:
forward: every node’s score, bottom-up (a composite is the sum of its children, a mixture is the log-sum-exp of
log_w_j + child_j);E-step backward: the responsibility reaching a node (the product of the mixture posteriors down its path, times the observation weight) is pushed to its children; each leaf accumulates its weighted sufficient statistic, and each mixture its per-component counts.
It reuses the LeafTemplate machinery by giving every leaf node a
(1,)-shaped parameter block indexed at k = 0 (the templates are written for [k] indexing).
Scope: scalar leaves (the common case for nested mixtures); a nested model containing a matrix / tabulated
/ categorical leaf returns None here and falls back to numpy. It is consulted only when the flat
analyze() declines, so the flat fast path is never perturbed.
- analyze_nested(model)[source]
Return (tree, ctx) for a nested scalar-leaf Composite/Mixture model, or None to fall back.
Restricted to genuinely nested models – depth-2 flat mixtures/composites are handled (faster) by the flat
analyze(), so this only fires when that one declines.
- fused_nested_seq_log_density(model, enc)[source]
- fused_nested_accumulate(model, enc, weights, return_ll=False)[source]