braintrace.compile_etrace_graph

braintrace.compile_etrace_graph#

braintrace.compile_etrace_graph(model, *model_args, include_hidden_perturb=True, include_recurrent_mixing=False, sparse_n=None, snap_max_jacobian_elements=16777216, 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 (brainstate.nn.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, optional) – Whether to include hidden perturbations in the graph. Default True.

  • include_recurrent_mixing (bool, optional) – 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).

  • sparse_n (int, optional) – SnAp order for recurrence_scope='sparse_n'. When given, every hidden group carries the n-step position neighbourhood derived from its own transition (implemented by the internal position_graph module) in HiddenGroup.snap, and its trace’s trailing state axis widens to K * num_state. Requires include_recurrent_mixing=True – the widened operator is gathered out of the coupled transition’s Jacobian – and raises rather than degrading to K = 1 if it is not set. None (default) for every other scope.

  • snap_max_jacobian_elements (int, optional) – Ceiling on each group’s widened block Jacobian, P * (K * S) ** 2 elements. Only consulted when sparse_n is given. Default DEFAULT_MAX_JACOBIAN_ELEMENTS.

  • control_flow (ControlFlowPolicy or None, optional) – 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 the internal scan_descent module;

    • 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;

    • caps every canonicalization fixpoint at fixpoint_iteration_limit sweeps (default 64) so control flow that never converges raises CompilationError naming the offending equations instead of hanging the compiler. Raise it for models that genuinely nest control flow deeper than that.

Returns:

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

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