mixle.inference.node_precision_plan module

Per-NODE precision planning for a composed distribution tree.

mixle.inference.precision_plan picks ONE compute precision for a whole model. This module generalizes that decision to every NODE of a composed tree (a MixtureDistribution of components, a CompositeDistribution of factors, and any nesting of the two): each node gets its own safety verdict, reusing the exact per-leaf safety check precision_plan already validates (family whitelist + variance floor), then those leaf verdicts are aggregated UP the tree. This is the roadmap’s “fits are deterministic given seed, and sufficient statistics are ADDITIVE, so error bounds compose like stats” insight applied literally: a non-leaf node is float32-safe iff every leaf beneath it is, and its advertised summed-LL error bound is the SUM of its leaves’ bounds (each leaf’s bound independently verified, see precision_plan’s module docstring).

Two things live here:

1. recommend_tree_precision() – walks the WHOLE tree and returns a TreePrecisionPlan: an inspectable, path-keyed mapping from every node (leaf and non-leaf) to its chosen precision and rationale. This is the “D1-reported property / D6-H3 action” surface: a caller (a future node-report or a block-freeze policy) can read exactly which sub-blocks are safe to run cheap, without re-deriving the verdict.

2. mixed_precision_fit() – actually EXECUTES an EM fit where each top-level child of the root combinator (each mixture component, or each composite factor) runs its E-step scoring and sufficient-statistic accumulation at ITS OWN assigned precision. See the “Execution scope” note in that function’s docstring for exactly how far genuine per-node execution reaches in the current architecture, and where it honestly falls back.

class NodePrecision(path, node_type, is_leaf, compute_dtype, rationale, rel_error_bound, leaf_count)

Bases: object

The precision verdict for ONE node of a composed tree.

Parameters:

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

• node_type (str)

• is_leaf (bool)

• compute_dtype (Any)

• rationale (str)

• rel_error_bound (float)

• leaf_count (int)

path

Field-path identifying this node from the tree root, e.g. ("components", "0", "dists", "1") for the second factor of the first mixture component. () is the root.

Type:

tuple[str, …]

node_type

The node’s class name ("MixtureDistribution", "CompositeDistribution", or the leaf’s own class name).

Type:

str

is_leaf

True for an actual distribution leaf (not a mixture/composite combinator).

Type:

bool

compute_dtype

The chosen dtype (np.float32 or np.float64).

Type:

Any

rationale

Human-readable reason, reusing precision_plan’s per-leaf wording where applicable.

Type:

str

rel_error_bound

The advertised relative summed-log-likelihood error bound FOR THIS NODE’s subtree (0.0 when compute_dtype is float64 – exact).

Type:

float

leaf_count

Number of leaves in this node’s subtree (1 for a leaf itself).

Type:

int

class TreePrecisionPlan(root_type, nodes=<factory>)

Bases: object

The full per-node precision plan for a composed tree: path -> NodePrecision.

This is the inspectable/actionable artifact the roadmap calls for: iterate nodes to see every block’s verdict, call dtype_for() to look up one node, or reduced_paths() / frozen_candidates() to drive a future block-freeze / precision-drop policy (D6/H3). Hooking this into D1’s NodeReport (once that lands) is a natural follow-up – see the module docstring.

Parameters:

• root_type (str)

• nodes (dict[tuple[str, ...], NodePrecision])

reduced_paths()

Paths (any node, leaf or subtree) allocated float32.

Return type:

list[tuple[str, …]]

top_level_child_paths()

Paths one level below the root – the granularity mixed_precision_fit() can actually execute at independently (see that function’s docstring for why).

Return type:

list[tuple[str, …]]

advertised_bound(path=())

The advertised relative summed-log-likelihood error bound for path’s subtree (default: whole tree). Sums the verified per-leaf bound (FUSED_FP32_REL_LL_BOUND) over every REDUCED leaf beneath path – the additive composition the roadmap calls for.

Parameters:

path (tuple[str, ...])

Return type:

float

recommend_tree_precision(model, data, min_variance=1e-6, max_magnitude=1e6, sample_size=4096)

Return the per-NODE precision plan for a composed tree.

Walks model (a Mixture / Composite / leaf, and any nesting thereof) and computes a safety verdict at EVERY node: leaves get the identical per-leaf check precision_plan uses (family whitelist + variance floor); non-leaf nodes aggregate their children (safe iff ALL children are safe) and their advertised error bound is the additive sum of their leaves’ bounds. The (global) data-magnitude check is evaluated once against data and gates every node uniformly, matching precision_plan.recommend_compute_precision – only the leaf/family conditioning genuinely varies node-to-node in this codebase (see this module’s and mixed_precision_fit’s docstrings for why data conditioning isn’t yet split per-node).

Parameters:

• model (Any) – The composed distribution tree (root).

• data (Any) – Representative data used to check the magnitude guard (see precision_plan).

• min_variance (float) – Leaves with sigma2 below this are treated as near-degenerate -> float64.

• max_magnitude (float) – Data magnitude guard, identical semantics to precision_plan.

• sample_size (int) – Stride-sample size used for the magnitude guard.

Returns:

A TreePrecisionPlan with one NodePrecision per node (root included, at path ()).

Return type:

TreePrecisionPlan

mixed_precision_fit(model, data, plan=None, max_its=10, delta=1.0e-9, weights=None)

Fit model with each TOP-LEVEL CHILD of the root combinator executing its E-step (scoring + sufficient-statistic accumulation) at its OWN precision, per plan.

Execution scope (read this before trusting “mixed precision” claims elsewhere): the numba fused kernel (mixle.stats.compute.fused_codegen) compiles ONE kernel per fusible subtree and runs it at ONE dtype end to end – there is no way to hand it two different literal dtypes inside a single call. So the finest granularity at which this codebase can genuinely execute different literal precisions within one fit is the boundary between independently-callable fused subtrees, which is exactly the immediate children of the root combinator: each mixture COMPONENT, or each composite FACTOR. Nesting deeper than that (e.g. two factors of a Composite that is itself one mixture component) shares one dtype – the whole subtree is one fused-kernel call, so it gets the AND-aggregated verdict recommend_tree_precision() already computes for it.

This is genuinely DIFFERENT from mixle.inference.optimize(precision=...): that entry point threads exactly one engine (one dtype) through the WHOLE fit via a single NumpyEngine / FusedKernel – there is currently no per-node engine plumbed through optimize’s EM loop. This function does NOT go through optimize; it is a standalone driver, scoped to a root MixtureDistribution or CompositeDistribution (any nesting below each top-level child is fine – it just shares that child’s one dtype, as described above). Anything else (a bare leaf, or a combinator this driver doesn’t recognize) is fit at plain float64 with a warning-free no-op fallback (there is nothing to split).

Each child accumulates sufficient statistics in its OWN accumulator, at its OWN dtype for the row arithmetic; every accumulator’s OUTPUT (and the softmax/logsumexp responsibility normalization for a mixture) is float64, matching the “accumulation is ALWAYS float64” invariant precision_plan documents – so, like the model-global allocator, results never drift regardless of which nodes ran reduced; only the per-row SCORE of the reduced nodes is computed cheaper.

Parameters:

• model (Any) – MixtureDistribution or CompositeDistribution to fit (used as both the shape AND the starting parameter estimate – pass an initialized model, e.g. from estimator().estimate or a previous optimize call).

• data (Any) – Training data.

• plan (TreePrecisionPlan | None) – A TreePrecisionPlan (e.g. from recommend_tree_precision()). None computes one internally against data.

• max_its (int) – Maximum EM iterations.

• delta (float | None) – Convergence threshold on the per-iteration total log-likelihood change. None runs exactly max_its iterations.

• weights (ndarray | None) – Optional per-observation weights (default: uniform 1.0).

Returns:

The fitted model (same top-level type as model).

Return type:

Any