HiddenPerturbation

class braintrace.HiddenPerturbation

The hidden-perturbation information.

Hidden perturbation adds a perturbation variable to each hidden state in the jaxpr and replaces the hidden states with the perturbed states:

\[h^t = f(x) \;\Rightarrow\; h^t = f(x) + \mathrm{perturb\_var},\]

where \(h\) is the hidden state, \(f\) is the function, \(x\) is the input, and \(\mathrm{perturb\_var}\) is the perturbation variable.

Variables:

• perturb_vars (sequence of Var) – The perturbation variables.

• perturb_hidden_paths (sequence of Path) – The hidden-state paths that are perturbed.

• perturb_hidden_states (sequence of brainstate.HiddenState) – The hidden states that are perturbed.

• perturb_jaxpr (ClosedJaxpr) – The perturbed jaxpr.

See also

add_hidden_perturbation_in_module

Build perturbations directly from a model.

Notes

Internally a new variable \(\hat{h}^t = f(x)\) is defined and an extra equation \(h^t = \hat{h}^t + \mathrm{perturb\_var}\) is added. The perturbation lets the hidden-state gradient be read off the perturbation variable

\[\frac{\partial L^t}{\partial h^t} = \frac{\partial L^t}{\partial \mathrm{perturb\_var}}.\]

Examples

>>> import brainstate
>>> import braintrace
>>> gru = braintrace.nn.GRUCell(3, 4)
>>> _ = brainstate.nn.init_all_states(gru)
>>> inputs = brainstate.random.randn(3)
>>> hidden_perturb = braintrace.add_hidden_perturbation_in_module(gru, inputs)
>>> isinstance(hidden_perturb, braintrace.HiddenPerturbation)
True

static __new__(_cls, perturb_vars, perturb_hidden_paths, perturb_hidden_states, perturb_jaxpr)

Create new instance of HiddenPerturbation(perturb_vars, perturb_hidden_paths, perturb_hidden_states, perturb_jaxpr)

dict()

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

Returns:

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

eval_jaxpr(inputs, perturb_data)

Evaluate the perturbed jaxpr.

Parameters:

• inputs (sequence of jax.Array) – The flat input values of the original jaxpr.

• perturb_data (sequence of jax.Array) – The perturbation values, one per entry of perturb_vars.

Returns:

sequence of jax.Array – The outputs of the perturbed jaxpr.

init_perturb_data()

Initialize the perturbation data to zeros.

Returns:

sequence of jax.Array – One zero array per perturbation variable, matching its shape and dtype.

perturb_data_to_hidden_group_data(perturb_data, hidden_groups)

Convert the perturbation data to per-hidden-group data.

Parameters:

• perturb_data (sequence of jax.Array) – The perturbation values, one per entry of perturb_vars.

• hidden_groups (sequence of HiddenGroup) – The hidden groups to map the perturbation data onto.

Returns:

sequence of jax.Array – One concatenated perturbation array per hidden group.

Raises:

ValueError – If perturb_data does not have the same length as perturb_vars, or if a hidden group needs a path that was not perturbed.

Notes

Both guards are if ... raise rather than assert so that python -O cannot strip them. The mapping from paths to perturbation data is positional (zip(perturb_hidden_paths, perturb_data)), so a length disagreement silently re-attributes every cotangent past the mismatch; and a group asking for an unperturbed path used to surface as a bare KeyError naming a path tuple and nothing else. Both checks read only Python-level lengths and dict keys, never array data.

perturb_hidden_paths

Alias for field number 1

perturb_hidden_states

Alias for field number 2

perturb_jaxpr

Alias for field number 3

perturb_vars

Alias for field number 0

HiddenPerturbation.__init__()