mixle.inference.node_precision_plan module

Executed-subtree 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). The proposed verdict is then reconciled with the subtrees the mixed-precision driver can actually execute. Reduced subtrees are compared against float64 on the supplied data and carry an absolute log-score enclosure. Relative likelihood errors are never added: cancellation can make that composition unbounded.

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, abs_log_error_bound, leaf_count, validation_rows=0, execution_scope='proposed')[source]

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)

  • abs_log_error_bound (float | None)

  • leaf_count (int)

  • validation_rows (int)

  • execution_scope (str)

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

abs_log_error_bound

Triangle-inequality enclosure on the absolute error of the unweighted summed log score over the validation rows. It is not a relative or out-of-sample bound.

Type:

float | None

leaf_count

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

Type:

int

validation_rows

Number of rows used to establish abs_log_error_bound.

Type:

int

execution_scope

Whether this node executes independently, inherits a parent’s dtype, or is combined in float64.

Type:

str

class TreePrecisionPlan(root_type, nodes=<factory>, validation_fingerprint=None, validation_rows=0)[source]

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

  • validation_fingerprint (str | None)

  • validation_rows (int)

reduced_paths()[source]

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

Return type:

list[tuple[str, …]]

top_level_child_paths()[source]

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=())[source]

Return the validation-data absolute summed-log-score error enclosure for path.

Parameters:

path (tuple[str, ...])

Return type:

float

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

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

Walks model (a Mixture / Composite / leaf, and any nesting thereof) and computes a safety proposed verdict at every node, then reconciles it with the independently callable subtrees the executor actually uses. Each reduced executable subtree is scored in float32 and float64 on data; the plan records a triangle-inequality absolute error enclosure. Descendants inherit their executable parent’s dtype and do not advertise independent bounds. The root mixture or composite combines child scores in float64.

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)[source]

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

Reduced precision applies only to scoring. Sufficient-statistic accumulation and mixture normalization remain float64. Every reduced score call is compared with float64 at runtime; a subtree that stops meeting the absolute error tolerance is permanently downgraded to float64 in the supplied plan.

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