ETraceVmap

class braintrace.ETraceVmap

Provide sequence drivers on a brainstate.nn.Vmap wrapper.

Returned by braintrace.compile(..., vmap=True) so the call site is identical in batched and unbatched mode. Because it is a brainstate.nn.Vmap, calling it, its attributes and every isinstance(x, brainstate.nn.Vmap) check keep working; only the added methods are new.

One thing does change: type(x) is brainstate.nn.Vmap is now False, so a caller dispatching on the exact runtime type takes a different branch, and repr reads ETraceVmap. Use isinstance. (Pickling is unaffected – a bare Vmap was already unpicklable, for the same weakref reason.)

Reaching into .module is not an equivalent: learner.module.etrace_grad(...) would drive the unbatched learner and silently produce per-lane-wrong results.

Window mode is refused here – see SequenceDriverMixin.etrace_grad().

ETraceVmap.__init__(module, in_axes=0, out_axes=0, vmap_states=None, vmap_out_states=None, axis_name=None, axis_size=None)