mixle.inference.backend_respecialization module

Backend re-specialization mid-fit – learned-scheduler-ready decision logic + execution machinery for swapping a NODE’s execution backend in response to structural changes observed during a fit (workstream D6).

Frame (see the ConditionalJIT track, D1-D6): the estimator tree is an IR. D1 (mixle.inference.node_report) instruments every node with an update_kind classification and an E/M cost proxy. D2 (mixle.inference.freeze_rollup) freezes subtrees that stop moving. D3 (mixle.inference.block_em) schedules WHICH blocks get updated each round, so a node’s “hot” (updated every round) vs. “cold” (rarely scheduled) status is a live, observable signal DURING a fit, not a static property. K1’s per-node precision plan (mixle.inference.precision_plan, and the wider per-node walk on the per-node-precision branch) reports when a node’s safe compute precision drops. Every one of these is a moment where the OPTIMAL execution BACKEND for a node – eager vs. torch.compile-d, full vs. reduced precision, computed-on-the-fly vs. table-cached – might also change.

Correctness backbone (unchanged from the rest of the D-track): compilation and density memoization are exact execution optimizations. Reduced precision is different and is exposed only when a separate validation supplies a finite absolute error bound. NodeBackend.apply() records only actions it actually executes and rejects unsupported fusion rather than pretending it occurred. Any interleaving of exact backend swaps with the block-EM schedule (D3) or freeze/roll-up caching (D2) remains coordinate ascent on the same Neal-Hinton free energy F.

Two things live here:

  1. Compile economics (estimate_compile_cost() / estimate_compile_benefit() / estimate_table_cost() / estimate_table_benefit(), and the trigger functions built on top of them) – a real cost/benefit tradeoff, not a fixed rule: an upfront re-specialization cost is only worth paying when it is amortized over enough EXPECTED remaining executions of the node at its current hot/frozen/precision status. RespecializationDecision is the inspectable receipt of that tradeoff – the “compile economics exposed to D5” interface the roadmap calls for: a later learned controller (D5, not built here) could plausibly consume its estimated_cost / estimated_benefit / net_benefit fields as training features and action as a label, without needing to re-derive the economics itself.

  2. Execution (NodeBackend, DensityTable, compile_forward()) – actually APPLIES a supported re-specialization: wraps a node’s forward call in torch.compile (reusing mixle.engines.torch_engine.TorchEngine.compile()’s exact convention – the SAME compile_enabled and hasattr(torch, "compile") gate, not a parallel one), or swaps in a state-aware exact density-table lookup for repeated identical queries. This is a real mechanism with a measurable effect, not just a recommendation (see the acceptance tests).

class RespecializationAction(*values)[source]

Bases: StrEnum

The backend choice a RespecializationDecision recommends (and NodeBackend can actually apply).

class RespecializationDecision(field_path, node_type, action, triggered_by, estimated_cost, estimated_benefit, expected_remaining_calls, rationale, validated_abs_error_bound=None, target_precision=None)[source]

Bases: object

The cost/benefit tradeoff and chosen action for ONE node’s backend – the D5-facing receipt.

This is deliberately a flat, inspectable dataclass (not just a boolean): estimated_cost, estimated_benefit, and expected_remaining_calls are exactly the features a later learned controller (D5, out of scope here) would want as training signal, and action / triggered_by double as the label and the (attributable) reason. Nothing in this dataclass is invented for D5’s benefit alone – every field is also what THIS module’s own decision functions compute and act on today.

Parameters:

  • field_path (str)

  • node_type (str)

  • action (RespecializationAction)

  • triggered_by (str)

  • estimated_cost (float)

  • estimated_benefit (float)

  • expected_remaining_calls (float)

  • rationale (str)

  • validated_abs_error_bound (float | None)

  • target_precision (Any | None)

property net_benefit: float

estimated_benefit - estimated_cost – positive iff re-specializing is worth it.

estimate_compile_cost(report, *, fixed_overhead=_DEFAULT_COMPILE_FIXED_OVERHEAD, per_param_overhead=_DEFAULT_COMPILE_PER_PARAM_OVERHEAD)[source]

Estimate the upfront cost of torch.compile-ing report’s node: a fixed graph-capture/tracing overhead plus a per-parameter term (more parameters -> a bigger graph to trace and specialize), mirroring D1’s own param_count-proxy convention rather than inventing a new cost unit.

Parameters:

  • report (NodeReport)

  • fixed_overhead (float)

  • per_param_overhead (float)

Return type:

float

estimate_compile_benefit(report, expected_remaining_calls, *, per_call_cost=None, speedup_factor=_DEFAULT_COMPILE_SPEEDUP_FACTOR)[source]

Estimate the total cost SAVED by compiling report’s node, amortized over expected_remaining_calls more executions at its current per-call cost (per_call_cost, or D1’s own E/M cost proxy when not supplied). A real cost-benefit tradeoff, not a fixed rule: a node executed many more times pays back a fixed compile overhead; a node executed only a handful more times does not (see the decision-boundary test in mixle.tests.backend_respecialization_test).

Parameters:

  • report (NodeReport)

  • expected_remaining_calls (float)

  • per_call_cost (float | None)

  • speedup_factor (float)

Return type:

float

estimate_table_cost(n_query_points, *, fixed_overhead=_DEFAULT_TABLE_FIXED_OVERHEAD, per_point_overhead=_DEFAULT_TABLE_PER_POINT_OVERHEAD)[source]

Estimate the upfront cost of building a density table over n_query_points seed points: a fixed bookkeeping overhead plus one evaluation per seed point.

Parameters:

  • n_query_points (int)

  • fixed_overhead (float)

  • per_point_overhead (float)

Return type:

float

estimate_table_benefit(report, expected_remaining_calls, *, per_call_cost=None, speedup_factor=_DEFAULT_TABLE_SPEEDUP_FACTOR)[source]

Estimate the total cost saved by serving report’s node from a precomputed density table instead of computing on the fly – same amortization shape as estimate_compile_benefit(), with a table’s own (typically larger, since a dict lookup is cheaper than a re-traced graph call) default speedup factor.

Parameters:

  • report (NodeReport)

  • expected_remaining_calls (float)

  • per_call_cost (float | None)

  • speedup_factor (float)

Return type:

float

decide_hot_compile(report, *, activation_ratio, expected_remaining_calls, already_compiled=False, hot_threshold=_DEFAULT_HOT_ACTIVATION_RATIO, per_call_cost=None, speedup_factor=_DEFAULT_COMPILE_SPEEDUP_FACTOR)[source]

Decide whether report’s node should be compiled because D3’s scheduler has been running it “hot” (active in activation_ratio fraction of recent rounds, >= hot_threshold).

A frozen node (D2/D1’s own update_kind) is never a compile candidate regardless of a stale activation ratio – it is not going to run again. An already-compiled node is left alone (no double-compile).

Parameters:

  • report (NodeReport)

  • activation_ratio (float)

  • expected_remaining_calls (float)

  • already_compiled (bool)

  • hot_threshold (float)

  • per_call_cost (float | None)

  • speedup_factor (float)

Return type:

RespecializationDecision

decide_frozen_precision_drop(report, *, q_gain_tol=_DEFAULT_STABLE_Q_GAIN_TOL, already_reduced=False, validated_abs_error_bound=None, target_precision=np.float32)[source]

Decide whether a frozen/converged node may use an independently validated lower precision.

Frozen status and small objective gain do not establish numerical safety. The action is therefore emitted only when validated_abs_error_bound is supplied by an executable precision comparison and is finite and non-negative.

Parameters:

  • report (NodeReport)

  • q_gain_tol (float)

  • already_reduced (bool)

  • validated_abs_error_bound (float | None)

  • target_precision (Any)

Return type:

RespecializationDecision

decide_density_table(report, *, expected_remaining_calls, n_query_points, structure_stable, per_call_cost=None, speedup_factor=_DEFAULT_TABLE_SPEEDUP_FACTOR)[source]

Decide whether a fully closed-form node (update_kind in {"closed_form", "conjugate_closed_form"} – no gradient loop, nothing to compile) whose surrounding tree structure has stabilized (structure_stable, e.g. D3’s scheduler has stopped changing which blocks are active) is worth precomputing a density table for.

Parameters:

  • report (NodeReport)

  • expected_remaining_calls (float)

  • n_query_points (int)

  • structure_stable (bool)

  • per_call_cost (float | None)

  • speedup_factor (float)

Return type:

RespecializationDecision

compile_forward(fn, engine=None)[source]

Wrap fn with torch.compile via a compile-enabled engine’s own .compile method (mixle.engines.torch_engine.TorchEngine.compile()) – reuses that method’s exact compile_enabled and hasattr(torch, "compile") gate so this module never has a second, possibly-divergent opinion about when compilation is available.

Parameters:
Return type:

Callable[[…], Any]

class DensityTable(fn, seed_points=None, *, quantum=None, state_token=None)[source]

Bases: object

An exact, state-aware memoization table for a closed-form node’s forward function.

Keys contain the input type, device, dtype, shape, and exact bytes; no quantization or near-neighbor approximation is performed. state_token must return a hashable signature of every mutable value read by fn. A token change clears the cache before lookup, preventing results computed under stale model parameters from being reused.

Parameters:

  • fn (Callable[[Any], Any])

  • seed_points (list[Any] | None)

  • quantum (float | None)

  • state_token (Callable[[], Any] | None)

invalidate()[source]

Discard all memoized results explicitly.

Return type:

None

lookup(x)[source]

Return an exact state-matched cache hit, or compute and store a detached snapshot.

Parameters:

x (Any)

Return type:

Any

class NodeBackend(dist, forward=None, *, engine=None, state_token=None)[source]

Bases: object

Holds and executes the CURRENTLY CHOSEN backend for one node’s forward call.

Wraps an eager forward callable (defaulting to dist.kernel(engine=...).score, the same engine-aware evaluation kernel the rest of the codebase already uses – see mixle.stats.compute.kernel) and lets apply() swap in a compiled or table-cached variant per a RespecializationDecision, while __call__ always dispatches to whichever backend is currently active. Compilation and density-table execution preserve the exact input semantics; reduced precision requires a separately validated error-bound receipt.

Parameters:

  • dist (Any)

  • forward (Callable[[Any], Any] | None)

  • engine (Any | None)

  • state_token (Callable[[], Any] | None)

apply(decision, *, table_seed_points=None)[source]

Apply a supported decision and record only the backend that actually became active.

Parameters:

  • decision (RespecializationDecision)

  • table_seed_points (list[Any] | None)

Return type:

None