ETraceAlgorithm

class braintrace.ETraceAlgorithm

Provide the base interface for eligibility-trace algorithms.

Parameters:

• model (brainstate.nn.Module) – The model function, which receives the input arguments and returns the model output.

• graph_executor (ETraceGraphExecutor) – The executor used to evaluate the compiled eligibility-trace graph.

• name (str, optional) – The name of the etrace algorithm.

Variables:

• graph (ETraceGraph) – The compiled eligibility-trace graph.

• executor (ETraceGraphExecutor) – The executor associated with graph.

• report (CompilationReport) – Structured diagnostics produced by graph compilation.

• param_states (Dict[Hashable, brainstate.ParamState]) – The weight states.

• hidden_states (Dict[Hashable, brainstate.HiddenState]) – The hidden states.

• other_states (Dict[Hashable, brainstate.State]) – The other states.

• is_compiled (bool) – Whether the etrace algorithm has been compiled.

• running_index (brainstate.LongTermState[int]) – Zero-based count of previously-completed update() calls: it reads 0 during the first call, 1 during the second, and so on – i.e. the value used inside a given update() is the pre-increment count, not the 1-based call number. It is incremented by one immediately after each update() completes, and reset back to 0 by reset_state().

compile_graph(*args)

Compile the eligibility trace graph of the relationship between etrace weights, states and operators.

The compilation process includes:

• building the etrace graph

• separating the states

• initializing the etrace states

Parameters:

*args – The input arguments.

property executor

Get the etrace graph executor.

Returns:

ETraceGraphExecutor – The etrace graph executor.

get_etrace_of(weight)

Get the eligibility trace of the given weight.

Parameters:

weight (brainstate.ParamState | Path) – The parameter weight or path to the weight.

Returns:

Any – The eligibility trace.

Raises:

NotImplementedError – This method must be implemented by subclasses.

property graph

Get the etrace graph.

Returns:

ETraceGraph – The etrace graph.

property hidden_states

Get the hidden states.

Returns:

brainstate.util.FlattedDict[Path, brainstate.HiddenState] – The hidden states.

init_etrace_state(*args, **kwargs)

Initialize the eligibility trace states of the etrace algorithm.

This method is needed after compiling the etrace graph. See .compile_graph() for the details.

Parameters:

• *args – The positional arguments.

• **kwargs – The keyword arguments.

Raises:

NotImplementedError – This method must be implemented by subclasses.

property other_states

Get the other states.

Returns:

brainstate.util.FlattedDict[Path, brainstate.State] – The other states.

property param_states

Get the parameter weight states.

Returns:

brainstate.util.FlattedDict[Path, brainstate.ParamState] – The parameter weight states.

property path_to_states

Get the path to the states.

Returns:

brainstate.util.FlattedDict[Path, brainstate.State] – The mapping from path to states.

property report

Structured, read-only report of the compiled eligibility-trace graph.

Returns:

CompilationReport – A view over graph.

Raises:

RuntimeError – If accessed before the graph has been compiled.

show_graph(verbose=True, return_msg=False)

Display the eligibility-trace graph.

Delegates to ETraceGraphExecutor.show_graph().

Parameters:

• verbose (bool, optional) – If True (default), print the summary to stdout.

• return_msg (bool, optional) – If True, return the summary string. Default False.

Returns:

None or str – The summary string if return_msg is True, else None.

property state_id_to_path

Get the state ID to the path.

Returns:

Dict[int, Path] – The mapping from state ID to path.

update(*args, **kwargs)

Update the model and the eligibility trace states.

Parameters:

• *args – The input arguments.

• **kwargs – Per-call options, not model inputs. See __call__.

Returns:

Any – The model output.

Raises:

NotImplementedError – This method must be implemented by subclasses.

ETraceAlgorithm.__init__(model, graph_executor, name=None)