ControlFlowPolicy

class braintrace.ControlFlowPolicy

Policy knobs governing control-flow canonicalization.

Parameters:

• cond (str, optional) – How ETP-relevant cond equations are handled. 'convert' (default) if-converts them into inlined branches + select_n; 'opaque' leaves every cond untouched, so the existing control-flow restrictions apply (weights used inside raise NotImplementedError; ETP primitives inside are excluded with a warning).

• scan_unroll_limit (int, optional) – Maximum static length of an ETP-relevant inner scan that unroll_inner_scans unrolls into flat equations. Longer (or effectful, or while-containing) scans stay opaque with a DiagnosticKind.SCAN_UNROLL_SKIPPED warning, and the existing control-flow restrictions apply. 0 (or negative) disables unrolling entirely. Default 16.

• while_hidden (str, optional) – How an opaque control-flow equation (while, or a scan/ cond the canonicalizer left opaque) that produces a hidden-state output without consuming any weight invar is handled. 'opaque-fwd' (default) keeps it as an opaque forward node: the whole equation is embedded in the hidden-to-hidden transition and its Jacobian is extracted in forward mode (while has no reverse-mode rule), recorded as a DiagnosticKind.CONTROL_FLOW_OPAQUE_FWD INFO diagnostic. 'error' restores the pre-opaque-forward behaviour of raising NotImplementedError.

• etp_in_control_flow (str, optional) – How an ETP primitive found inside a remaining opaque control-flow body is handled during relation discovery. 'error' (default) raises NotImplementedError — that weight would otherwise silently drop out of online learning. 'exclude' restores the previous behaviour of a loud warning (DiagnosticKind.PRIMITIVE_INSIDE_CONTROL_FLOW) plus exclusion of the weight from ETP relations.

• scan_descent (str, optional) – How ETP-relevant scans too long to unroll are handled. 'auto' (default) rewrites the scan for structured descent: relations and hidden groups are discovered inside the body and the eligibility trace is folded over the substep axis (see braintrace._compiler.scan_descent). 'off' preserves the pre-Phase-4 behavior: the scan stays opaque and compilation fails on the existing control-flow restrictions.

• fixpoint_iteration_limit (int, optional) – Maximum number of sweeps a canonicalization fixpoint may run before giving up with a CompilationError naming the equations the last sweep was still rewriting. Default 64.

  This bounds loop iterations, not the size of any single rewrite — the per-rewrite bound is scan_unroll_limit. Each sweep rewrites the currently visible cond/scan equations and re-inlines the jit bodies they surface, so the number of sweeps a jaxpr needs is set by how deeply its ETP-relevant control flow nests, plus one final sweep that rewrites nothing and so proves the fixpoint was reached. Real models nest a handful of levels deep; raise the limit if yours genuinely nests deeper. There is deliberately no “unbounded” setting: without a cap, a jaxpr that regenerates convertible control flow as fast as the sweeps consume it hangs the compiler with no diagnostic at all. Must be a positive integer.

Notes

If-conversion changes execution semantics: both branches of a converted cond execute every step, and select_n discards the dead branch’s value. Guarded partial operations (a cond used to avoid sqrt of a negative number, for example) therefore need care:

• Values and forward-mode (JVP) derivatives are safe — select_n selects whole outputs and tangents, so a dead-branch NaN/Inf never reaches the selected result.

• Reverse-mode (VJP) gradients are NOT safe when the dead branch’s local Jacobian is NaN/Inf: select_n’s transpose hands the dead branch an exact-zero cotangent, but the dead branch’s VJP multiplies that zero by its NaN/Inf Jacobian (0 * nan = nan), contaminating gradients of inputs shared with the live branch — the classic single-where pitfall. If a cond guards a partial operation’s domain, keep it opaque (cond='opaque') or guard the operand itself (e.g. sqrt(where(ok, x, 1.))) inside the branch.

For branches whose values and Jacobians are finite, the canonicalized jaxpr is exact in both value and derivative. Branches with effects, containing while, or containing a scan that unroll_inner_scans could not unroll, are never converted (see if_convert_conds). Scan unrolling itself has no semantics change: the unrolled equations compute exactly what the loop computed.

Examples

>>> import braintrace
>>> policy = braintrace.ControlFlowPolicy(cond='opaque')
>>> policy.cond
'opaque'

ControlFlowPolicy.__init__(cond='convert', scan_unroll_limit=16, while_hidden='opaque-fwd', etp_in_control_flow='error', scan_descent='auto', fixpoint_iteration_limit=64)