HiddenGroup

class braintrace.HiddenGroup(index: int, hidden_paths: List[Tuple[str, ...]], hidden_states: List[HiddenState], hidden_invars: List[Var], hidden_outvars: List[Var], transition_jaxpr: Jaxpr, transition_jaxpr_constvars: List[Var], is_diagonal_recurrence: bool = True, descent: GroupDescent | None = None)

The data structure recording a hidden-group relation.

A hidden group bundles the hidden states that are mutually connected through a recurrence transition, together with the jaxpr that computes that transition

\[h_1^t, h_2^t, \ldots = f(h_1^{t-1}, h_2^{t-1}, \ldots, x^t).\]

index

Position of this group in the compiled group sequence.

Type:

int

hidden_paths

The module path to each hidden state in the group.

Type:

list of Path

hidden_states

The hidden states in the group.

Type:

list of brainstate.HiddenState

hidden_invars

The input jaxpr Var of each hidden state (at the previous step).

Type:

list of HiddenInVar

hidden_outvars

The output jaxpr Var of each hidden state (at the current step).

Type:

list of HiddenOutVar

transition_jaxpr

The jaxpr computing the hidden-state transition for the group.

Type:

Jaxpr

transition_jaxpr_constvars

The other input variables required to evaluate transition_jaxpr.

Type:

list of Var

is_diagonal_recurrence

Whether the recurrence is diagonal across the leading varshape positions (see the field comment for the full contract).

Type:

bool

descent

Descent context when this group’s transition is one substep of a descended scan (Phase 4 structured scan descent); None for ordinary groups.

Type:

GroupDescent or None

See also

find_hidden_groups_from_module

Build hidden groups directly 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)
>>> hidden_groups, _ = braintrace.find_hidden_groups_from_module(gru, inputs)
>>> len(hidden_groups)
1

check_consistent_varshape()

Check whether the shapes of the hidden states are consistent.

Raises:

NotSupportedError – If the shapes of the hidden states are not consistent.

concat_hidden(splitted_hid_vals)

Concatenate split hidden-state values into a single array.

Concatenates a sequence of split hidden-state values along the last axis. For non-HiddenGroupState values, an extra trailing dimension is added before concatenation.

Parameters:

splitted_hid_vals (Sequence[Array]) – A sequence of split hidden-state values, each corresponding to a hidden state in the group.

Returns:

A single array containing all hidden-state values concatenated along the last axis.

Return type:

jax.Array

descent: GroupDescent | None

Alias for field number 8

diagonal_jacobian(hidden_vals, input_vals)

Compute the diagonal Jacobian matrix along the last dimension.

Parameters:

• hidden_vals (Sequence[Array]) – The hidden-state values.

• input_vals (Any) – The input values.

Returns:

The per-position block-diagonal of the recurrent Jacobian d h^t / d h^{t-1}, with shape (*varshape, num_states, num_states). Entry [p, a, b] is d h^t[p, a] / d h^{t-1}[p, b] – the cross-position terms d h^t[p] / d h^{t-1}[q] (p != q) are intentionally dropped (the D-RTRL / e-prop diagonal approximation).

Return type:

jax.Array

Notes

For diagonal recurrence (is_diagonal_recurrence is True) the positions are independent, so the cheap column-sum produced by jacrev_last_dim() already equals this block diagonal. For coupled recurrence the column sum would instead add in the off-diagonal cross-position terms – inflating every entry and driving the eligibility trace to overflow – so the true block diagonal is extracted directly via block_diagonal_last_dim().

When the transition contains a while equation (an opaque forward node), reverse-mode differentiation is unavailable (JAX has no transpose rule for while), so the Jacobian is extracted in forward mode instead (jacfwd_last_dim(), or block_diagonal_last_dim() with use_forward_mode=True) – same values, different derivative mode.

dict()

Return this group’s named fields as a plain dictionary.

Returns:

An ordered mapping from field name to value, as produced by the underlying typing.NamedTuple.

Return type:

Dict[str, Any]

hidden_invars: List[Var]

Alias for field number 3

hidden_outvars: List[Var]

Alias for field number 4

hidden_paths: List[Tuple[str, ...]]

Alias for field number 1

hidden_states: List[HiddenState]

Alias for field number 2

index: int

Alias for field number 0

is_diagonal_recurrence: bool

Alias for field number 7

property num_state: int

The number of hidden states.

Returns:

The total number of hidden states across the group.

Return type:

int

split_hidden(concat_hid_vals)

Split a concatenated hidden-state array into individual arrays.

Splits a concatenated array of hidden-state values into separate arrays, one per hidden state in the group. HiddenGroupState and non-HiddenGroupState values are handled differently.

Parameters:

concat_hid_vals (Array) – A concatenated array of hidden-state values. The last dimension is assumed to contain the concatenated states.

Returns:

A list of split hidden-state arrays. For non-HiddenGroupState values, the last dimension is squeezed.

Return type:

list of jax.Array

transition(hidden_vals, input_vals)

Compute the hidden-state transitions.

Evaluates the group transition jaxpr

\[h_1^t, h_2^t, \cdots = f(h_1^{t-1}, h_2^{t-1}, \cdots, x^t).\]

Parameters:

• hidden_vals (Sequence[Array]) – The old hidden-state values.

• input_vals (Any) – The input values.

Returns:

The new hidden-state values.

Return type:

List[Array]

transition_jaxpr: Jaxpr

Alias for field number 5

transition_jaxpr_constvars: List[Var]

Alias for field number 6

property varshape: Tuple[int, ...]

The shape of each state variable.

Returns:

The variable shape shared by the hidden states in the group.

Return type:

tuple of int