SequenceDriverMixin

class braintrace.SequenceDriverMixin

etrace_grad / etrace_evolve, written once for both hosts.

Hosts supply three hooks:

_seq_call

The callable that drives one step or window. Defaults to self.

_seq_param_states

The default set of brainstate.ParamState to differentiate.

_seq_vjp_method

The learner’s vjp_method, or None if it has none.

_seq_vjp_method is a hook rather than a getattr on the driver because brainstate.nn.Vmap defines no __getattr__ and so does not forward vjp_method from .module. Reading the attribute off the driver object would silently yield None for every vmapped learner and bypass the window-mode validation entirely.

etrace_evolve(*sequences, step_fn=None, chunk_size=None, return_outputs=False)

Drive the model and the eligibility trace forward, computing no gradient.

Hidden states and eligibility traces advance exactly as they do inside etrace_grad().

Parameters:

• *sequences – As in etrace_grad().

• step_fn (callable, optional) – None (default) calls the learner directly with the slices – and, under window mode, wraps them in MultiStepData, since with no step_fn every sequence is by definition a model input. Supplying a step_fn opts out of that wrapping entirely.

• chunk_size (int, optional) – As in etrace_grad(), except that windows are legal on either vjp_method: this method runs no loss VJP, which is the only thing single-step learners refuse.

• return_outputs (bool, optional) – False (default) returns None and stacks nothing, so a long warm-up costs no output memory. True stacks whatever the call returned, with leading axis T // chunk_size.

Returns:

None or stacked outputs – None when return_outputs=False; otherwise the stacked per-step (or per-window) return values of the driven call.

Raises:

• TypeError – As in etrace_grad(), for chunk_size and for wrapped sequences.

• ValueError – As in etrace_grad(), except that a window is not refused for being on a single-step learner – no loss VJP is taken here, so the restriction does not apply. Windows are still refused under a vmapped learner, and T % chunk_size == 0 still holds.

Examples

>>> import braintrace
>>> learner = braintrace.compile(model, 'D_RTRL', xs[0], batch_size=1)
>>> warmup_inputs, xs, ys = inputs[:20], inputs[20:], targets[20:]
>>>
>>> learner.etrace_evolve(warmup_inputs)          # free-running prefix
>>> grads = learner.etrace_grad(xs, ys, step_fn=step_loss)

etrace_grad(*sequences, step_fn, mask=None, chunk_size=None, weights=None, reduction='mean', loss_output='per_step', has_aux=False, return_value=False)

Accumulate online gradients over a sequence.

Parameters:

• *sequences – One or more pytrees whose leaves share a leading length T, sliced in lockstep and passed to step_fn positionally. There is no distinguished targets argument. May not be a SingleStepData / MultiStepData wrapper – wrap inside step_fn instead.

• step_fn (callable) – The user’s step function, which runs the model itself and returns the loss. Keyword-only and required. It must call this learner exactly once per invocation: zero calls leave the trace un-advanced for that step, two advance it twice, and neither is detectable from the returned gradient.

• mask (array, optional) – (T,) per-step loss weights; None means all-ones. Values need not be binary. Gates only the loss – the model and the eligibility trace are driven at every step regardless.

• chunk_size (int, optional) – None (default) or 1 drive step-by-step, handing seq[t] to step_fn, which returns a scalar. k >= 2 drives in windows, handing seq[t:t+k] of shape (k, ...), and step_fn must return a (k,) vector of per-step losses and wrap its model inputs in MultiStepData. Window mode requires vjp_method='multi-step' and T % k == 0, and is not available under a vmapped learner.

• weights (dict, optional) – The brainstate.ParamState to differentiate. Defaults to the learner’s own param_states.

• reduction ({‘mean’, ‘sum’}, optional) – 'mean' (default) divides by the total mask weight, max(sum_t mask_t, 1), not by T.

• loss_output ({‘per_step’, ‘masked’, ‘scalar’}, optional) – What return_value=True hands back: the raw pre-mask losses (T,), the masked losses (T,), or the reduced objective (scalar). Ignored when return_value=False.

• has_aux (bool, optional) – Whether step_fn returns (loss, aux).

• return_value (bool, optional) – Whether to return the losses alongside the gradients.

Returns:

grads or tuple – Mirrors brainstate.transform.grad: grads, (grads, losses), (grads, aux) or (grads, losses, aux).

Raises:

• TypeError – If step_fn is not given; if chunk_size is not None or a Python int (a traced or numpy value is refused, since the value has to be known at trace time); or if a sequence is a SingleStepData / MultiStepData wrapper. The wrappers are registered pytree nodes, so slicing one would decompose the wrapper rather than the data.

• ValueError – If no sequences are given, if their leading lengths disagree, or if T == 0 – there is nothing to slice. If chunk_size is below 1; if k >= 2 but the learner’s vjp_method is not 'multi-step' (the executor would raise three frames down), the learner is vmapped (in_axes=0 would map time as the batch axis), or T % k != 0. If mask is not shape (T,). If reduction or loss_output is not one of its legal values. If step_fn returns a non-scalar in plain mode, or anything but shape (k,) in window mode. If the learner has not been compiled, the existing guard raises before any of these.

Warning

Two known limitations are reachable through this method. Neither is a driver defect – both belong to the engines it drives – but chunk_size is what makes them easy to hit by accident, so they are named here rather than only in the limitations document.

F-30, window mode on an IO-factorized engine. pp_prop / ES_D_RTRL / OSTLFeedforward, and any trace_factorization='io_factorized' config, undo the warm-up bias of their f-side smoothing with a factor indexed by running_index, which counts update() calls. A k-step window advances the trace k times but running_index once, so the correction lags the trace by a factor of k. Measured at T=6, k=2, decay=0.9: 2.1e-03 relative against the same run indexed by true trace steps. Driving with chunk_size=None avoids it entirely.

F-35, DNI and the synthesizer’s deployment contract. A train_synthetic_gradient() fit is valid only for the exact (loss_fn, chunk_size) pair it was trained on, and nothing checks that deployment matches. Fit and drive at the same chunk_size – note that 1 and None are the same path on both sides. A mismatch is not degraded DNI but noise shaped like a cotangent: fitting at 1 and deploying at 2 moves the gradient 2.0e-02 relative. mask is the other half of the contract, since it changes the objective the synthesizer is predicting the future of. reduction is not: it divides once after the loop and never enters the differentiated objective.

Notes

grads is the learner’s online-gradient estimate of the reduced objective, not in general its mathematical derivative. For an exact algorithm inside its valid regime the two coincide; for every approximate rule they deliberately do not, and that difference is the algorithm’s content. reduction and mask define the objective the estimate is aimed at.

Examples

>>> import braintools
>>> import braintrace
>>> learner = braintrace.compile(model, 'D_RTRL', inputs[0], batch_size=1)
>>> opt = braintools.optim.Adam(1e-3)
>>> opt.register_trainable_weights(learner.param_states)
>>>
>>> def step_loss(inp, tar):
...     out = learner(inp)
...     return braintools.metric.squared_error(out, tar).mean()
>>>
>>> grads, loss = learner.etrace_grad(
...     inputs, targets, step_fn=step_loss,
...     loss_output='scalar', return_value=True)
>>> opt.update(grads)

SequenceDriverMixin.__init__()