braintrace.compile

Contents

braintrace.compile#

braintrace.compile(model, algorithm, *example_inputs, batch_size=None, seed=None, verbose=0, vmap=False, **options)#

Define an eligibility-trace online-learning model in one call.

This is the unified entry point. It initializes the model’s states, builds the eligibility-trace graph, checks that the model is trainable online, and optionally prints a compilation report before returning a ready-to-update learner.

Parameters:

  • model (brainstate.nn.Module) – The recurrent / spiking model defining one-step behavior. It does not need to be pre-initialized; compile always (re)initializes its states.

  • algorithm (type, str or ETraceConfig) – An ETraceAlgorithm subclass, a registered case-insensitive name, or an ETraceConfig naming the learning-rule coordinate directly. Registered names are 'd_rtrl', 'pp_prop', 'es_d_rtrl', 'esd_rtrl', 'eprop', 'e_prop', 'ostl_recurrent', 'ostl_feedforward', 'uoro', 'three_factor', and 'dni'. The aliases 'es_d_rtrl' and 'esd_rtrl' select pp_prop; 'e_prop' selects EProp. SnAp has no registered string name and must be passed as a class. A config selects the engine through its trace_factorization and is forwarded to the constructor, so any coordinate admitted by the compatibility matrix can be compiled without a named preset.

  • *example_inputs (Any) – Example call inputs (arrays / SingleStepData / MultiStepData) matching what learner.update(...) will receive. At least one is required.

  • batch_size (int or None, optional) – Forwarded to brainstate.nn.init_all_states. None (default) initializes unbatched states. Must match the batch dimension of example_inputs.

  • seed (int or None, optional) – If given, state initialization runs inside brainstate.random.seed_context for reproducibility; the global RNG is restored afterwards. None (default) leaves the RNG untouched. Weights created at model-construction time are outside this scope.

  • verbose (int, optional) – Report verbosity printed at compile time: 0 (default) silent, 1 the structural summary, 2 additionally compiler WARNING/ERROR diagnostics. Other values raise ValueError.

  • vmap (bool, optional) – When False (default) states are initialized with init_all_states(model, batch_size=batch_size). When True, states are created under brainstate.transform.vmap_new_states(state_tag='new', axis_size=batch_size) and the learner is wrapped in ETraceVmap. In vmap mode: example_inputs carry the batch axis (axis 0); batch_size is required and used as the vmap axis_size; the return value is a ETraceVmap whose .module is the unbatched learner (use result.module.report for its report). Drive sequences through the returned wrapper, never through result.module. Requires a model whose hidden states are all (re)created in init_all_states; models holding construction-time states may raise brainstate.transform.BatchAxisError.

  • **options (Any) – Forwarded to the algorithm constructor. See Algorithm options below.

Returns:

ETraceAlgorithm or ETraceVmap – When vmap=False, the compiled learner carries a report; call .update(*inputs) to train. When vmap=True, returns an ETraceVmap wrapper (also a brainstate.nn.Vmap); access the underlying learner’s report as .module.report. Call etrace_grad and etrace_evolve on the wrapper itself, not on .module.

Raises:

  • ValueError – If algorithm is an unknown name, no example_inputs are given, verbose is not in {0, 1, 2}, or vmap=True without batch_size.

  • TypeError – If algorithm is not an ETraceAlgorithm subclass, registered string, or ETraceConfig; if a required algorithm option is missing; or if the same config is supplied both as algorithm and through config=.

  • braintrace.CompilationError – If no trainable weights are routed through ETP ops (nothing to learn online).

Notes

Algorithm options. **options are forwarded verbatim to the algorithm constructor. Required options have no default; omitting one raises TypeError. The class pages are the authoritative source for accepted options:

Passing an algorithm class supports subclasses and avoids the string registry. Passing ETraceConfig selects ParamDimVjpAlgorithm, IODimVjpAlgorithm, or RandomProjectionVjpAlgorithm from trace_factorization. Do not also pass config= in that case.

Calling compile twice on the same model re-initializes its states.

Axis coordinates. Passing an ETraceConfig in the algorithm position compiles a coordinate that may have no preset name:

# an x-side leak with an instantaneous f-side
learner = braintrace.compile(
    model,
    braintrace.ETraceConfig(
        trace_factorization='io_factorized', decay=(0.9, 0.0)),
    x0,
)

Examples

>>> import brainstate
>>> import braintrace
>>> import jax.numpy as jnp
>>>
>>> class RNN(brainstate.nn.Module):
...     def __init__(self):
...         super().__init__()
...         self.cell = braintrace.nn.ValinaRNNCell(3, 4, activation='tanh')
...         self.out = braintrace.nn.Linear(4, 1)
...     def update(self, x):
...         return x >> self.cell >> self.out
>>>
>>> model = RNN()
>>> x0 = brainstate.random.randn(1, 3)   # (batch, features)
>>> # Registered string: initialize states, build the graph, return a learner.
>>> by_name = braintrace.compile(model, 'd_rtrl', x0, batch_size=1)
>>> y = by_name.update(x0)
>>>
>>> # Algorithm classes, including SnAp, can be passed directly.
>>> by_class = braintrace.compile(
...     model, braintrace.D_RTRL, x0, batch_size=1)
>>>
>>> # A config selects an engine from its trace factorization.
>>> config = braintrace.ETraceConfig()
>>> by_config = braintrace.compile(
...     model, config, x0, batch_size=1)
>>>
>>> # Every learner returned by compile carries the two sequence drivers.
>>> xs = brainstate.random.randn(10, 1, 3)   # (T, batch, features)
>>> ys = brainstate.random.randn(10, 1, 1)
>>> def step_loss(x, y):
...     return jnp.mean((by_name(x) - y) ** 2)
>>> grads, losses = by_name.etrace_grad(xs, ys, step_fn=step_loss, return_value=True)
>>> # etrace_evolve is the same drive with no loss: it advances hidden
>>> # state and the eligibility trace, optionally stacking the outputs.
>>> outs = by_name.etrace_evolve(xs, return_outputs=True)
Contents