ETraceGraph

ETraceGraph#

class braintrace.ETraceGraph#

The overall compiled graph for the eligibility trace.

Tracks the relationship between the eligibility-trace weights (ParamState), the eligibility-trace variables (HiddenState), and the eligibility-trace operations (ETP primitives). It is the object returned by compile_etrace_graph() and consumed by the online-learning algorithms.

Variables:

  • module_info (ModuleInfo) – The model information.

  • hidden_groups (sequence of HiddenGroup) – The hidden groups.

  • hid_path_to_group (dict) – Mapping from each hidden-state path to its HiddenGroup.

  • hidden_param_op_relations (sequence of HiddenParamOpRelation) – The hidden parameter-operation relations.

  • hidden_perturb (HiddenPerturbation or None) – The hidden perturbation, or None when perturbations are excluded.

  • diagnostics (tuple of CompilationRecord) – The structured compilation records emitted while building the graph.

See also

compile_etrace_graph

Build an ETraceGraph from a model.

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)
>>> isinstance(graph, braintrace.ETraceGraph)
True
static __new__(_cls, module_info, hidden_groups, hid_path_to_group, hidden_param_op_relations, hidden_perturb, diagnostics=())#

Create new instance of ETraceGraph(module_info, hidden_groups, hid_path_to_group, hidden_param_op_relations, hidden_perturb, diagnostics)

call_hidden_perturb(args, perturb_data, old_state_vals=None)#

Run the forward pass with additive perturbations injected at the hidden states.

Evaluates the perturbed-forward jaxpr built during compilation, which is the forward computation augmented so that each tracked hidden state has a perturbation term added to it. This is the primitive used to probe hidden->hidden and hidden->output sensitivities.

Parameters:

  • args (Inputs) – The model inputs for this step, matching the signature captured at compile time.

  • perturb_data (Sequence[jax.Array]) – One perturbation array per tracked hidden state, added at the corresponding perturbation site.

  • old_state_vals (Sequence[jax.Array] or None, optional) – The state values to run from. When None (default) the current values of the compiled model states are used.

Returns:

tuple – (outputs, etrace_state_vals, other_state_vals, temp_data) – the same four-element structure produced by a normal forward call through ModuleInfo.jaxpr_call().

diagnostics#

Alias for field number 5

dict()#

Return the graph’s fields as a plain dictionary.

Returns:

dict – A mapping from field name to value for every attribute of this ETraceGraph.

explain(*, weight_path=None, hidden_path=None, kind=None)#

Return compilation records filtered by weight path, hidden path, or kind.

weight_path and hidden_path match the record’s weight_path exactly and hidden_paths membership respectively. kind matches CompilationRecord.kind. All filters are optional; with no filters the full diagnostic log is returned.

Parameters:

  • weight_path (Path or None, optional) – If given, keep only records whose weight_path equals this value. Default None.

  • hidden_path (Path or None, optional) – If given, keep only records whose hidden_paths contain this value. Default None.

  • kind (DiagnosticKind or None, optional) – If given, keep only records whose kind is this value. Default None.

Returns:

tuple of CompilationRecord – The matching records, in emission order.

hid_path_to_group#

Alias for field number 2

hidden_groups#

Alias for field number 1

hidden_param_op_relations#

Alias for field number 3

hidden_perturb#

Alias for field number 4

module_info#

Alias for field number 0

ETraceGraph.__init__()#