ETraceAlgorithm

class braintrace.ETraceAlgorithm(model, graph_executor, name=None)

The base class for the eligibility trace algorithm.

Parameters:

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

• name (str | None) – The name of the etrace algorithm.

graph

The etrace graph.

Type:

ETraceGraphExecutor

param_states

The weight states.

Type:

Dict[Hashable, brainstate.ParamState]

hidden_states

The hidden states.

Type:

Dict[Hashable, brainstate.HiddenState]

other_states

The other states.

Type:

Dict[Hashable, brainstate.State]

is_compiled

Whether the etrace algorithm has been compiled.

Type:

bool

running_index

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().

Type:

brainstate.LongTermState[int]

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 (Any) – The input arguments.

Return type:

None

property executor: ETraceGraphExecutor

Get the etrace graph executor.

Returns:

The etrace graph executor.

Return type:

ETraceGraphExecutor

get_etrace_of(weight)

Get the eligibility trace of the given weight.

Parameters:

weight (ParamState | Tuple[str, ...]) – The parameter weight or path to the weight.

Returns:

The eligibility trace.

Return type:

Any

Raises:

NotImplementedError – This method must be implemented by subclasses.

property graph: ETraceGraph

Get the etrace graph.

Returns:

The etrace graph.

Return type:

ETraceGraph

property hidden_states: FlattedDict

Get the hidden states.

Returns:

The hidden states.

Return type:

brainstate.util.FlattedDict[Path, brainstate.HiddenState]

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 (Any) – The positional arguments.

• **kwargs (Any) – The keyword arguments.

Raises:

NotImplementedError – This method must be implemented by subclasses.

Return type:

None

property other_states: FlattedDict

Get the other states.

Returns:

The other states.

Return type:

brainstate.util.FlattedDict[Path, brainstate.State]

property param_states: FlattedDict

Get the parameter weight states.

Returns:

The parameter weight states.

Return type:

brainstate.util.FlattedDict[Path, brainstate.ParamState]

property path_to_states: FlattedDict

Get the path to the states.

Returns:

The mapping from path to states.

Return type:

brainstate.util.FlattedDict[Path, brainstate.State]

property report: CompilationReport

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

Returns:

A view over graph.

Return type:

CompilationReport

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) – If True (default), print the summary to stdout.

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

Returns:

The summary string if return_msg is True, else None.

Return type:

None | str

property state_id_to_path: Dict[int, Tuple[str, ...]]

Get the state ID to the path.

Returns:

The mapping from state ID to path.

Return type:

Dict[int, Path]

update(*args)

Update the model and the eligibility trace states.

Parameters:

*args (Any) – The input arguments.

Returns:

The model output.

Return type:

Any

Raises:

NotImplementedError – This method must be implemented by subclasses.