compile_etrace_graph

class braintrace.compile_etrace_graph(model, *model_args, include_hidden_perturb=True, include_recurrent_mixing=False, control_flow=None)

Construct the eligibility-trace graph for a given model and inputs.

This is the primary entry point of the ETrace compiler. It builds the graph for the model, tracking the relationship between the eligibility-trace weights (ParamState), the eligibility-trace states (HiddenState), and the eligibility-trace operations (ETP primitives). These relationships are later used to compute the weight spatial gradients, the hidden-state Jacobian, and the hidden-state-to-weight Jacobian.

Parameters:

• model (Module) – The model for which the eligibility-trace graph is built.

• *model_args (Tuple) – The positional arguments required by the model.

• include_hidden_perturb (bool) – Whether to include hidden perturbations in the graph. Default True.

• include_recurrent_mixing (bool) – Hidden-group grouping mode for the hidden-to-hidden transition. When False (default, “without recurrence”), recurrent ETP mixing primitives (e.g. the recurrent etp_mv/etp_mm) are treated as boundaries and excluded from the transition jaxpr, so the transition is element-wise and the per-position recurrent Jacobian is diagonal (the bounded D-RTRL / e-prop approximation). When True (“with recurrence”), those primitives are traced into the transition, the recurrence becomes coupled, and the true per-position block-diagonal Jacobian is extracted (RTRL-exact temporal credit, e.g. for OSTLRecurrent / OSTTP).

• control_flow (ControlFlowPolicy | None) –

  Policy governing control-flow canonicalization and downstream handling, forwarded to extract_module_info() and (via ModuleInfo.control_flow) to every later compiler pass. None (default) uses the default policy, which:

  • if-converts every ETP-relevant cond into inlined branches + select_n (both branches then execute every step); pass ControlFlowPolicy(cond='opaque') to restore the previous behavior (weights inside cond raise NotImplementedError);

  • unrolls every ETP-relevant scan of static length at most scan_unroll_limit (default 16);

  • applies structured scan descent (scan_descent='auto') to ETP-relevant scans above the unroll limit: relations and hidden groups are discovered inside the scan body, the equation is rewritten to emit stacked per-substep values as extra ys, and the eligibility trace is folded over the substep axis at runtime — compile size stays independent of the scan length. Pass ControlFlowPolicy(scan_descent='off') to restore the pre-Phase-4 error. See braintrace._compiler.scan_descent;

  • keeps a weight-free while that reads/updates hidden state as an opaque forward node (while_hidden='opaque-fwd'): hidden-to-hidden Jacobians for groups whose transition crosses the loop are extracted in forward mode, and the perturbation pass detaches the loop’s inputs with stop_gradient so the perturbed jaxpr stays reverse-traceable. Pass ControlFlowPolicy(while_hidden='error') to reject such loops instead. A weight used through an ETP primitive inside a while is always a hard error;

  • raises on ETP primitives left inside a control-flow body the canonicalizer could not flatten (etp_in_control_flow='error'); pass ControlFlowPolicy(etp_in_control_flow='exclude') to restore the warn-and-exclude behavior.

Returns:

The compiled eligibility-trace graph containing module information, hidden groups, hidden parameter-operation relations, and optional hidden perturbations.

Return type:

ETraceGraph

Raises:

NotImplementedError – If a recursive call to the compiler is detected.

See also

ETraceGraph

The returned compiled-graph data structure.

Examples

>>> import brainstate
>>> import braintrace
>>> gru = braintrace.nn.GRUCell(3, 4)
>>> _ = brainstate.nn.init_all_states(gru)
>>> inputs = brainstate.random.randn(3)
>>> graph = braintrace.compile_etrace_graph(gru, inputs)
>>> len(graph.hidden_groups)
1