HiddenGroup

class braintrace.HiddenGroup

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).\]

Variables:

• index (int) – Position of this group in the compiled group sequence.

• hidden_paths (list of Path) – The module path to each hidden state in the group.

• hidden_states (list of brainstate.HiddenState) – The hidden states in the group.

• hidden_invars (list of HiddenInVar) – The input jaxpr Var of each hidden state (at the previous step).

• hidden_outvars (list of HiddenOutVar) – The output jaxpr Var of each hidden state (at the current step).

• transition_jaxpr (Jaxpr) – The jaxpr computing the hidden-state transition for the group.

• transition_jaxpr_constvars (list of Var) – The other input variables required to evaluate transition_jaxpr.

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

• snap (SnapPattern or None) – The SnAp-n neighbourhood the trace is widened onto (recurrence_scope='sparse_n'); None for every other scope.

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

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

static __new__(_cls, index, hidden_paths, hidden_states, hidden_invars, hidden_outvars, transition_jaxpr, transition_jaxpr_constvars, is_diagonal_recurrence=True, snap=None, descent=None)

Create new instance of HiddenGroup(index, hidden_paths, hidden_states, hidden_invars, hidden_outvars, transition_jaxpr, transition_jaxpr_constvars, is_diagonal_recurrence, snap, descent)

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 of jax.Array) – A sequence of split hidden-state values, each corresponding to a hidden state in the group.

Returns:

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

Raises:

ValueError – If splitted_hid_vals does not have exactly one entry per hidden state in the group.

Notes

The length check is not decorative. Before it existed this method zipped the value list against hidden_states, and zip truncates to the shorter argument: a short value list produced a concatenated array with a too-narrow trailing axis instead of an error, so a mis-routed cotangent surfaced later as a shape mismatch in unrelated trace math – or not at all, when the widths happened to coincide. It is an if ... raise and not an assert so that python -O cannot strip it.

The check reads only Python-level lengths, never array data, so it costs nothing in the compiled program.

descent

Alias for field number 9

diagonal_jacobian(hidden_vals, input_vals)

Compute the diagonal Jacobian matrix along the last dimension.

Parameters:

• hidden_vals (sequence of jax.Array) – The hidden-state values.

• input_vals (PyTree) – The input values.

Returns:

jax.Array – The per-position block-diagonal of the recurrent Jacobian d h^t / d h^{t-1}, with shape (*varshape, num_states, num_states) – or, when snap is set, the SnAp-n widened operator of shape (*varshape, trace_state_width, trace_state_width) whose entry [p, (k, a), (k', b)] is d h^t[nbr[p,k], a] / d h^{t-1}[nbr[p,k'], b]. 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).

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:

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

full_jacobian(hidden_vals, input_vals)

Compute the complete within-group hidden-to-hidden Jacobian.

The sibling of diagonal_jacobian() that keeps the cross-position terms that method drops. Selected by the graph executor when the algorithm’s trace_factorization is 'random_projection': UORO’s rank-1 estimator is unbiased for the recursion it rolls, so rolling the block diagonal would make it an unbiased estimate of an already-biased trace (matrix rule 11 rejects that coordinate).

Parameters:

• hidden_vals (sequence of jax.Array) – The hidden-state values.

• input_vals (PyTree) – The input values.

Returns:

jax.Array – Shape (*varshape, num_state, *varshape, num_state), with entry [p, a, q, b] = d h^t[p, a] / d h^{t-1}[q, b].

Notes

snap is ignored: the SnAp neighbourhood is a sparsity pattern for a stored trace, and this Jacobian is consumed immediately by a matrix-vector product rather than stored, so there is nothing to sparsify. Rule 11 rejects recurrence_scope='sparse_n' under 'random_projection' for that reason, so the combination cannot reach here.

The Jacobian is only full if the recurrent ETP mixing was traced into the transition, i.e. under include_recurrent_mixing. Rule 11 guarantees that by requiring recurrence_scope='coupled'.

hidden_invars

Alias for field number 3

hidden_outvars

Alias for field number 4

hidden_paths

Alias for field number 1

hidden_states

Alias for field number 2

index

Alias for field number 0

is_diagonal_recurrence

Alias for field number 7

property num_state

The number of hidden states.

Returns:

int – The total number of hidden states across the group.

snap

Alias for field number 8

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 (jax.Array) – A concatenated array of hidden-state values. The last dimension is assumed to contain the concatenated states.

Returns:

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

Raises:

ValueError – If the trailing axis of concat_hid_vals is not num_state wide, i.e. the array is not this group’s concatenated slab.

Notes

The width check is the inverse of concat_hidden()’s length check and closes the same silent path. u.math.split at this group’s cumulative boundaries returns one part per hidden state plus a trailing remainder, and that remainder is dropped by the zip below – so a too wide slab silently loses its surplus, and a too narrow one silently yields empty parts. Like the sibling check this is an if ... raise rather than an assert so python -O cannot strip it, and it reads only the static shape.

property trace_state_width

The width of a trace leaf’s trailing axis.

num_state normally; K * num_state under SnAp-n, where the trailing axis carries a (neighbour, state) pair rather than a state alone. Every trace allocation, recursion and solve reads this rather than num_state, which keeps the per-primitive kernels – all generic in that axis’s size – untouched.

Returns:

int – The trailing-axis width of this group’s trace leaves.

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 of jax.Array) – The old hidden-state values.

• input_vals (PyTree) – The input values.

Returns:

list of jax.Array – The new hidden-state values.

transition_jaxpr

Alias for field number 5

transition_jaxpr_constvars

Alias for field number 6

property varshape

The shape of each state variable.

Returns:

tuple of int – The variable shape shared by the hidden states in the group.

HiddenGroup.__init__()