ETraceGraph

class braintrace.ETraceGraph(module_info: ModuleInfo, hidden_groups: Sequence[HiddenGroup], hid_path_to_group: Dict[Tuple[str, ...], HiddenGroup], hidden_param_op_relations: Sequence[HiddenParamOpRelation], hidden_perturb: HiddenPerturbation | None, diagnostics: Tuple[CompilationRecord, ...] = ())

The overall compiled graph for the eligibility trace.

The eligibility trace graph, tracking the relationship between the etrace weights ParamState, the etrace variables HiddenState, and the etrace operations :pyETP primitives.

The following fields are included:

• module_info: The model information, instance of ModuleInfo.

• hidden_groups: The hidden groups, sequence of HiddenGroup.

• hid_path_to_group: The mapping from the hidden path to the hidden group HiddenGroup.

• hidden_param_op_relations: The hidden parameter operation relations, sequence of HiddenParamOpRelation.

• hidden_perturb: The hidden perturbation, instance of HiddenPerturbation, or None.

Example:

>>> import braintrace
>>> import brainstate
>>> gru = braintrace.nn.GRUCell(10, 20)
>>> gru.init_state()
>>> inputs = brainstate.random.randn(10)
>>> compiled_graph = braintrace.compile_etrace_graph(gru, inputs)
>>> compiled_graph.dict().keys()

diagnostics: Tuple[CompilationRecord, ...]

Alias for field number 5

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 are optional; with no filters the full diagnostic log is returned.

Return type:

Tuple[CompilationRecord, ...]

hid_path_to_group: Dict[Tuple[str, ...], HiddenGroup]

Alias for field number 2

hidden_groups: Sequence[HiddenGroup]

Alias for field number 1

hidden_param_op_relations: Sequence[HiddenParamOpRelation]

Alias for field number 3

hidden_perturb: HiddenPerturbation | None

Alias for field number 4

module_info: ModuleInfo

Alias for field number 0