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): re-specialization is a SCHEDULING/
EXECUTION optimization only. It never changes what a node computes – NodeBackend.__call__()
must return the tolerance-equal value regardless of which backend is currently active (see
mixle.tests.backend_respecialization_test’s tolerance-equal test) – only how fast it computes
it. Any interleaving of these backend swaps with the block-EM schedule (D3) or freeze/roll-up
caching (D2) is still coordinate ascent on the SAME Neal-Hinton free energy F the rest of the
track climbs; F itself never notices which backend answered a query.
Two things live here:
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.RespecializationDecisionis 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 itsestimated_cost/estimated_benefit/net_benefitfields as training features andactionas a label, without needing to re-derive the economics itself.Execution (
NodeBackend,DensityTable,compile_forward()) – actually APPLIES a chosen re-specialization: wraps a node’s forward call intorch.compile(reusingmixle.engines.torch_engine.TorchEngine.compile()’s exact convention – the SAMEcompile_enabled and hasattr(torch, "compile")gate, not a parallel one), or swaps in a precomputed density-table lookup for repeated identical/near-identical queries. This is a real mechanism with a measurable effect, not just a recommendation (see the acceptance tests).
- class RespecializationAction(*values)[source]
-
The backend choice a
RespecializationDecisionrecommends (andNodeBackendcan actually apply).
- class RespecializationDecision(field_path, node_type, action, triggered_by, estimated_cost, estimated_benefit, expected_remaining_calls, rationale)[source]
Bases:
objectThe 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, andexpected_remaining_callsare exactly the features a later learned controller (D5, out of scope here) would want as training signal, andaction/triggered_bydouble 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:
- 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-ingreport’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 ownparam_count-proxy convention rather than inventing a new cost unit.
- 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 overexpected_remaining_callsmore 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 inmixle.tests.backend_respecialization_test).
- 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_pointsseed points: a fixed bookkeeping overhead plus one evaluation per seed point.
- 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 asestimate_compile_benefit(), with a table’s own (typically larger, since a dict lookup is cheaper than a re-traced graph call) default speedup factor.
- 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 inactivation_ratiofraction 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).
- decide_frozen_precision_drop(report, *, q_gain_tol=_DEFAULT_STABLE_Q_GAIN_TOL, already_reduced=False)[source]
Decide whether a newly-frozen/near-converged node (D2’s freeze signal, or D1’s own near-zero Q-gain) should drop to reduced precision – the K1 “precision drops” trigger. Unlike compile/table decisions this one has no amortization term: a frozen node’s remaining work is (by definition) near zero, so the ONLY real cost is the negligible act of re-tagging its compute dtype, and the benefit is every future (typically read-only/health-check) touch of the node running cheaper – so this is a near-free action whenever the trigger fires at all.
- 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_kindin{"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.
- compile_forward(fn, engine=None)[source]
Wrap
fnwithtorch.compilevia a compile-enabled engine’s own.compilemethod (mixle.engines.torch_engine.TorchEngine.compile()) – reuses that method’s exactcompile_enabled and hasattr(torch, "compile")gate so this module never has a second, possibly-divergent opinion about when compilation is available.
- class DensityTable(fn, seed_points=None, *, quantum=1.0e-9)[source]
Bases:
objectA precomputed cache of a closed-form node’s density/forward function, keyed by a quantized input – a real memoization mechanism for repeated identical/near-identical query patterns (see the D6 module docstring).
quantumcontrols how near “near-identical” is: inputs whose quantized representation matches an existing key are served from the table; anything else falls back to (and populates the table with) the underlying function.
- class NodeBackend(dist, forward=None, *, engine=None)[source]
Bases:
objectHolds and executes the CURRENTLY CHOSEN backend for one node’s forward call.
Wraps an eager
forwardcallable (defaulting todist.kernel(engine=...).score, the same engine-aware evaluation kernel the rest of the codebase already uses – seemixle.stats.compute.kernel) and letsapply()swap in a compiled or table-cached variant per aRespecializationDecision, while__call__always dispatches to whichever backend is currently active. Every backend must return the tolerance-equal value for the same input – re-specializing changes only HOW the value is computed (see the module docstring’s correctness backbone).- Parameters:
dist (Any)
forward (Callable[[Any], Any] | None)
engine (Any | None)