ETraceGraphExecutor#
- class braintrace.ETraceGraphExecutor(model)#
The eligibility trace graph executor.
This class is used for computing the weight spatial gradients and the hidden state residuals. It is the most foundational class for the ETrace algorithms.
It is important to note that the graph is built no matter whether the model is batched or not. This means that this graph can be applied to any kind of models. However, the compilation is sensitive to the shape of hidden states.
- Parameters:
model (
Module) – The model to build the eligibility trace graph. The models should only define the one-step behavior.
- compile_graph(*args)[source]#
Build the eligibility trace graph for the model based on the provided inputs.
This method is crucial for constructing the graph used in the eligibility trace algorithm, which is essential for calculating weight spatial gradients and the hidden state Jacobian.
- Parameters:
*args – Positional arguments for the model, which may include inputs, parameters, or other necessary data required for graph compilation.
- Returns:
This method does not return any value. It initializes the compiled graph attribute of the instance.
- Return type:
- property graph: ETraceGraph#
Retrieve the compiled eligibility trace graph for the model.
This property provides access to the compiled graph, which is a crucial data structure for the eligibility trace algorithm. It contains various attributes that describe the relationships between the model’s variables, states, and operations.
- Returns:
The compiled graph for the model. This graph includes detailed information about the model’s structure, such as output variables, state variables, hidden-to-hidden variable relationships, and more.
- Return type:
- Raises:
ValueError – If the graph has not been compiled yet. Ensure to call the
compile_graph()` method before accessing this property –
- property path_to_states: FlattedDict[Tuple[str, ...], State]#
The path to the states.
- Returns:
The path to the states.
- show_graph(verbose=True, return_msg=False)[source]#
Display the graph illustrating the relationships between weights, operators, and hidden states.
This function generates a detailed message that describes the structure of the graph, including hidden groups, dynamic states, and weight parameters associated with hidden states. It can either print this message to the console or return it as a string.
- Parameters:
- Returns:
If return_msg is True, returns a string containing the graph details. Otherwise, returns None.
- Return type:
- solve_h2w_h2h_jacobian(*args)[source]#
Compute the hidden-to-weight and hidden-to-hidden Jacobian matrices.
This function is designed to calculate the forward propagation of the hidden-to-weight Jacobian and the hidden-to-hidden Jacobian based on the provided inputs and parameters. It is a crucial part of the eligibility trace algorithm, which helps in understanding the influence of weights and previous hidden states on the current hidden state.
Now we mathematically define what computations are done in this function.
For the state transition function \(y, h^t = f(h^{t-1}, \theta, x)\), this function aims to solve:
The function output: \(y\)
The updated hidden states: \(h^t\)
3. The Jacobian matrix of hidden-to-weight, i.e., \(\partial h^t / \partial \theta^t\). 2. The Jacobian matrix of hidden-to-hidden, i.e., \(\partial h^t / \partial h^{t-1}\).
- Parameters:
*args – Positional arguments for the model, which may include inputs, parameters, or other necessary data required for the computation of the Jacobians.
- Returns:
A tuple containing the following elements: - The function output (e.g., model predictions). - The updated hidden states after the current computation step. - Other states that may be relevant to the model’s operation. - The spatial gradients of the weights, represented by the hidden-to-weight Jacobian.
- Return type:
- solve_h2w_h2h_l2h_jacobian(*args)[source]#
Compute the hidden-to-weight and hidden-to-hidden Jacobian matrices, along with the VJP transformed loss-to-hidden gradients based on the provided inputs.
This function is designed to calculate both the forward propagation of the hidden-to-weight Jacobian and the loss-to-hidden gradients at the current time-step. It is essential for understanding the influence of weights and previous hidden states on the current hidden state, as well as the impact of the loss on the hidden states.
Particularly, this function aims to solve:
The Jacobian matrix of hidden-to-weight. That is, \(\partial h / \partial w\), where \(h\) is the hidden state and \(w\) is the weight.
The Jacobian matrix of hidden-to-hidden. That is, \(\partial h / \partial h\), where \(h\) is the hidden state.
The partial gradients of the loss with respect to the hidden states. That is, \(\partial L / \partial h\), where \(L\) is the loss and \(h\) is the hidden state.
- Parameters:
*args – Positional arguments for the model, which may include inputs, parameters, or other necessary data required for the computation of the Jacobians and gradients.
- Returns:
A tuple containing the following elements: - The function output (e.g., model predictions). - The updated hidden states after the current computation step. - Other states that may be relevant to the model’s operation. - The spatial gradients of the weights, represented by the hidden-to-weight Jacobian. - The residuals, which are the partial gradients of the loss with respect to the hidden states.
- Return type: