diff --git a/braintrace/_algorithm/base.py b/braintrace/_algorithm/base.py index 94d62f32..b403164e 100644 --- a/braintrace/_algorithm/base.py +++ b/braintrace/_algorithm/base.py @@ -37,9 +37,16 @@ class EligibilityTrace(brainstate.ShortTermState): - """ - The state for storing the eligibility trace during the computation of - online learning algorithms. + """Store the eligibility trace carried by an online-learning algorithm. + + Parameters + ---------- + value : PyTree[ArrayLike] + The initial eligibility-trace value. + name : str, optional + The state name. + **metadata : Any + Additional metadata stored with the state. Examples -------- diff --git a/braintrace/_algorithm/sequence.py b/braintrace/_algorithm/sequence.py index b7880518..1b1d2fbc 100644 --- a/braintrace/_algorithm/sequence.py +++ b/braintrace/_algorithm/sequence.py @@ -162,7 +162,7 @@ class SequenceDriverMixin: 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 :class:`brainstate.nn.Vmap` defines no ``__getattr__`` and so does + 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. @@ -285,7 +285,7 @@ def etrace_grad( Returns ------- grads or tuple - Mirrors :func:`brainstate.transform.grad`: ``grads``, + Mirrors ``brainstate.transform.grad``: ``grads``, ``(grads, losses)``, ``(grads, aux)`` or ``(grads, losses, aux)``. Raises @@ -537,7 +537,7 @@ def body(*slices): class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): - """A :class:`brainstate.nn.Vmap` that carries the sequence drivers. + """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 @@ -555,7 +555,8 @@ class ETraceVmap(SequenceDriverMixin, brainstate.nn.Vmap): would drive the **unbatched** learner and silently produce per-lane-wrong results. - Window mode is refused here -- see :meth:`SequenceDriverMixin._seq_check_window`. + Window mode is refused here -- see + :meth:`SequenceDriverMixin.etrace_grad`. """ __module__ = 'braintrace' diff --git a/braintrace/_compile.py b/braintrace/_compile.py index 0b4dfc16..cc96a19a 100644 --- a/braintrace/_compile.py +++ b/braintrace/_compile.py @@ -154,7 +154,7 @@ def compile( ``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 + *example_inputs : Any Example call inputs (arrays / :class:`SingleStepData` / :class:`MultiStepData`) matching what ``learner.update(...)`` will receive. At least one is required. @@ -184,7 +184,7 @@ def compile( ``result.module.report``). 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 + **options : Any Forwarded to the algorithm constructor. See *Algorithm options* below. Returns diff --git a/docs/_templates/classtemplate.rst b/docs/_templates/classtemplate.rst index eeb823a9..0bc710e3 100644 --- a/docs/_templates/classtemplate.rst +++ b/docs/_templates/classtemplate.rst @@ -7,3 +7,6 @@ .. autoclass:: {{ name }} :members: + :exclude-members: __init__ + +.. automethod:: {{ name }}.__init__ diff --git a/docs/apis/algorithm_details/braintrace.DNI.rst b/docs/apis/algorithm_details/braintrace.DNI.rst deleted file mode 100644 index e13cb6c9..00000000 --- a/docs/apis/algorithm_details/braintrace.DNI.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -DNI -=== - -.. autoclass:: DNI(model, synthesizer=None, name=None, vjp_method='multi-step', fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of, attach_synthesizer, group_signal_shapes diff --git a/docs/apis/algorithm_details/braintrace.D_RTRL.rst b/docs/apis/algorithm_details/braintrace.D_RTRL.rst deleted file mode 100644 index 83fae8bf..00000000 --- a/docs/apis/algorithm_details/braintrace.D_RTRL.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -D_RTRL -====== - -.. autoclass:: D_RTRL(model, name=None, vjp_method='single-step', fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.EProp.rst b/docs/apis/algorithm_details/braintrace.EProp.rst deleted file mode 100644 index 950cb210..00000000 --- a/docs/apis/algorithm_details/braintrace.EProp.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -EProp -===== - -.. autoclass:: EProp(model, feedback='symmetric', kappa_filter_decay=0.0, random_feedback_key=None, name=None, vjp_method='single-step', fast_solve=True, **kwargs) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.ETraceAlgorithm.rst b/docs/apis/algorithm_details/braintrace.ETraceAlgorithm.rst deleted file mode 100644 index 5f943b54..00000000 --- a/docs/apis/algorithm_details/braintrace.ETraceAlgorithm.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ETraceAlgorithm -=============== - -.. autoclass:: ETraceAlgorithm(model, graph_executor, name=None) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.ETraceConfig.rst b/docs/apis/algorithm_details/braintrace.ETraceConfig.rst deleted file mode 100644 index 58cfd54d..00000000 --- a/docs/apis/algorithm_details/braintrace.ETraceConfig.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ETraceConfig -============ - -.. autoclass:: ETraceConfig - :members: diff --git a/docs/apis/algorithm_details/braintrace.ETraceVjpAlgorithm.rst b/docs/apis/algorithm_details/braintrace.ETraceVjpAlgorithm.rst deleted file mode 100644 index 524c15d2..00000000 --- a/docs/apis/algorithm_details/braintrace.ETraceVjpAlgorithm.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ETraceVjpAlgorithm -================== - -.. autoclass:: ETraceVjpAlgorithm(model, name=None, vjp_method='single-step', control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.ETraceVmap.rst b/docs/apis/algorithm_details/braintrace.ETraceVmap.rst deleted file mode 100644 index f3bde022..00000000 --- a/docs/apis/algorithm_details/braintrace.ETraceVmap.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ETraceVmap -========== - -.. autoclass:: ETraceVmap(module, in_axes=0, out_axes=0, vmap_states=None, vmap_out_states=None, axis_name=None, axis_size=None) - :members: etrace_grad, etrace_evolve, update, reset_state, init_state diff --git a/docs/apis/algorithm_details/braintrace.EligibilityTrace.rst b/docs/apis/algorithm_details/braintrace.EligibilityTrace.rst deleted file mode 100644 index a186510b..00000000 --- a/docs/apis/algorithm_details/braintrace.EligibilityTrace.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -EligibilityTrace -================ - -.. autoclass:: EligibilityTrace - :members: diff --git a/docs/apis/algorithm_details/braintrace.FixedRandomFeedback.rst b/docs/apis/algorithm_details/braintrace.FixedRandomFeedback.rst deleted file mode 100644 index d1661e36..00000000 --- a/docs/apis/algorithm_details/braintrace.FixedRandomFeedback.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -FixedRandomFeedback -=================== - -.. autoclass:: FixedRandomFeedback - :members: diff --git a/docs/apis/algorithm_details/braintrace.IODimVjpAlgorithm.rst b/docs/apis/algorithm_details/braintrace.IODimVjpAlgorithm.rst deleted file mode 100644 index 060f62df..00000000 --- a/docs/apis/algorithm_details/braintrace.IODimVjpAlgorithm.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -IODimVjpAlgorithm -================= - -.. autoclass:: IODimVjpAlgorithm(model, decay_or_rank, name=None, vjp_method='single-step', fast_solve=True, control_flow=None, config=None, random_feedback_key=None) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.KappaFilter.rst b/docs/apis/algorithm_details/braintrace.KappaFilter.rst deleted file mode 100644 index 12420399..00000000 --- a/docs/apis/algorithm_details/braintrace.KappaFilter.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -KappaFilter -=========== - -.. autoclass:: KappaFilter - :members: diff --git a/docs/apis/algorithm_details/braintrace.OSTLFeedforward.rst b/docs/apis/algorithm_details/braintrace.OSTLFeedforward.rst deleted file mode 100644 index c2953a91..00000000 --- a/docs/apis/algorithm_details/braintrace.OSTLFeedforward.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -OSTLFeedforward -=============== - -.. autoclass:: OSTLFeedforward(model, decay_or_rank=1e-06, name=None, **kwargs) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.OSTLRecurrent.rst b/docs/apis/algorithm_details/braintrace.OSTLRecurrent.rst deleted file mode 100644 index 94cc6569..00000000 --- a/docs/apis/algorithm_details/braintrace.OSTLRecurrent.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -OSTLRecurrent -============= - -.. autoclass:: OSTLRecurrent(model, name=None, vjp_method='single-step', fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.ParamDimVjpAlgorithm.rst b/docs/apis/algorithm_details/braintrace.ParamDimVjpAlgorithm.rst deleted file mode 100644 index 3b8be5f4..00000000 --- a/docs/apis/algorithm_details/braintrace.ParamDimVjpAlgorithm.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ParamDimVjpAlgorithm -==================== - -.. autoclass:: ParamDimVjpAlgorithm(model, name=None, vjp_method='single-step', fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, config=None, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.RandomProjectionVjpAlgorithm.rst b/docs/apis/algorithm_details/braintrace.RandomProjectionVjpAlgorithm.rst deleted file mode 100644 index 92dcd300..00000000 --- a/docs/apis/algorithm_details/braintrace.RandomProjectionVjpAlgorithm.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -RandomProjectionVjpAlgorithm -============================ - -.. autoclass:: RandomProjectionVjpAlgorithm(model, name=None, vjp_method='multi-step', fast_solve=True, control_flow=None, config=None, projection_key=42, projection_eps=1e-12, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.SequenceDriverMixin.rst b/docs/apis/algorithm_details/braintrace.SequenceDriverMixin.rst deleted file mode 100644 index b278ac32..00000000 --- a/docs/apis/algorithm_details/braintrace.SequenceDriverMixin.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -SequenceDriverMixin -=================== - -.. autoclass:: SequenceDriverMixin - :members: etrace_grad, etrace_evolve diff --git a/docs/apis/algorithm_details/braintrace.SnAp.rst b/docs/apis/algorithm_details/braintrace.SnAp.rst deleted file mode 100644 index c4829c34..00000000 --- a/docs/apis/algorithm_details/braintrace.SnAp.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -SnAp -==== - -.. autoclass:: SnAp(model, n=2, name=None, vjp_method='single-step', fast_solve=True, snap_max_jacobian_elements=16777216, **kwargs) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.SyntheticGradient.rst b/docs/apis/algorithm_details/braintrace.SyntheticGradient.rst deleted file mode 100644 index 378abc95..00000000 --- a/docs/apis/algorithm_details/braintrace.SyntheticGradient.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -SyntheticGradient -================= - -.. autoclass:: SyntheticGradient(group_shapes, hidden_width=None, scale=0.0, seed=0) - :members: param_values, states_dict, apply diff --git a/docs/apis/algorithm_details/braintrace.ThreeFactor.rst b/docs/apis/algorithm_details/braintrace.ThreeFactor.rst deleted file mode 100644 index 69fa4685..00000000 --- a/docs/apis/algorithm_details/braintrace.ThreeFactor.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -ThreeFactor -=========== - -.. autoclass:: ThreeFactor(model, name=None, vjp_method='single-step', modulator=None, fast_solve=True, trace_dtype=None, chunked_trace=True, control_flow=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.UORO.rst b/docs/apis/algorithm_details/braintrace.UORO.rst deleted file mode 100644 index ed07e431..00000000 --- a/docs/apis/algorithm_details/braintrace.UORO.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -UORO -==== - -.. autoclass:: UORO(model, name=None, vjp_method='multi-step', fast_solve=True, control_flow=None, projection_key=42, projection_eps=1e-12, random_feedback_key=None, snap_max_jacobian_elements=16777216) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.compile.rst b/docs/apis/algorithm_details/braintrace.compile.rst deleted file mode 100644 index 0806a191..00000000 --- a/docs/apis/algorithm_details/braintrace.compile.rst +++ /dev/null @@ -1,9 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -compile -======== - -.. autofunction:: compile diff --git a/docs/apis/algorithm_details/braintrace.pp_prop.rst b/docs/apis/algorithm_details/braintrace.pp_prop.rst deleted file mode 100644 index d7850d13..00000000 --- a/docs/apis/algorithm_details/braintrace.pp_prop.rst +++ /dev/null @@ -1,10 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -pp_prop -======== - -.. autoclass:: pp_prop(model, decay_or_rank, name=None, vjp_method='single-step', fast_solve=True, control_flow=None, config=None, random_feedback_key=None) - :members: graph, report, compile_graph, show_graph, update, init_etrace_state, reset_state, get_etrace_of diff --git a/docs/apis/algorithm_details/braintrace.train_synthetic_gradient.rst b/docs/apis/algorithm_details/braintrace.train_synthetic_gradient.rst deleted file mode 100644 index 1abfd90e..00000000 --- a/docs/apis/algorithm_details/braintrace.train_synthetic_gradient.rst +++ /dev/null @@ -1,9 +0,0 @@ -.. role:: hidden - :class: hidden-section -.. currentmodule:: braintrace - - -train_synthetic_gradient -======================== - -.. autofunction:: train_synthetic_gradient \ No newline at end of file diff --git a/docs/apis/algorithms.rst b/docs/apis/algorithms.rst index 6b4a96b8..824448f9 100644 --- a/docs/apis/algorithms.rst +++ b/docs/apis/algorithms.rst @@ -7,33 +7,6 @@ Online-Learning Algorithms :local: :depth: 1 -.. toctree:: - :hidden: - - algorithm_details/braintrace.compile - algorithm_details/braintrace.ETraceConfig - algorithm_details/braintrace.ETraceAlgorithm - algorithm_details/braintrace.SequenceDriverMixin - algorithm_details/braintrace.ETraceVmap - algorithm_details/braintrace.ETraceVjpAlgorithm - algorithm_details/braintrace.EligibilityTrace - algorithm_details/braintrace.ParamDimVjpAlgorithm - algorithm_details/braintrace.D_RTRL - algorithm_details/braintrace.IODimVjpAlgorithm - algorithm_details/braintrace.pp_prop - algorithm_details/braintrace.SnAp - algorithm_details/braintrace.RandomProjectionVjpAlgorithm - algorithm_details/braintrace.UORO - algorithm_details/braintrace.ThreeFactor - algorithm_details/braintrace.DNI - algorithm_details/braintrace.SyntheticGradient - algorithm_details/braintrace.train_synthetic_gradient - algorithm_details/braintrace.EProp - algorithm_details/braintrace.OSTLRecurrent - algorithm_details/braintrace.OSTLFeedforward - algorithm_details/braintrace.FixedRandomFeedback - algorithm_details/braintrace.KappaFilter - ``braintrace`` provides online-learning algorithms based on eligibility-trace propagation. They all share one interface: wrap a model, compile its graph, then call the learner as a drop-in replacement for the model's forward pass — @@ -53,6 +26,7 @@ for a model and eagerly builds its eligibility-trace graph, returning a ready-to-``update`` learner in a single call. .. autosummary:: + :toctree: generated/ :nosignatures: compile @@ -76,7 +50,7 @@ online learning used to require is not something you write by hand: inputs[n_warmup:], targets, step_fn=step_loss, return_value=True) opt.update(grads) -:meth:`~ETraceAlgorithm.etrace_grad` owns the loop, the accumulation, the loss +:meth:`~SequenceDriverMixin.etrace_grad` owns the loop, the accumulation, the loss mask and the reduction; ``step_fn`` owns the model call. That split is what lets a multi-head model, a hidden-state regularizer, or a windowed objective work without the driver knowing anything about them. Both methods are @@ -90,7 +64,9 @@ step and requires ``vjp_method='multi-step'``; ``chunk_size=1`` is the plain single-step path, matching :func:`train_synthetic_gradient`'s encoding. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst SequenceDriverMixin ETraceVmap @@ -139,7 +115,9 @@ Pass a config wherever :func:`compile` accepts an algorithm name: ) .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst ETraceConfig @@ -147,17 +125,23 @@ Pass a config wherever :func:`compile` accepts an algorithm name: Base Classes ------------ -The abstract bases shared by every algorithm. :class:`ETraceAlgorithm` is the -root; :class:`ETraceVjpAlgorithm` adds the VJP-based machinery that the -concrete D-RTRL / ES-D-RTRL / SNN algorithms build on. :class:`EligibilityTrace` -is the state these algorithms carry across time. +The abstract bases and reusable estimator engines shared across algorithms. +:class:`ETraceAlgorithm` is the root, :class:`ETraceVjpAlgorithm` adds the +VJP-based machinery, and :class:`EligibilityTrace` is the state carried across +time. The three estimator bases implement the parameter-dimensional, +input/output-factorized, and random-projection trace representations. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst ETraceAlgorithm ETraceVjpAlgorithm EligibilityTrace + ParamDimVjpAlgorithm + IODimVjpAlgorithm + RandomProjectionVjpAlgorithm D-RTRL — Parameter-dimensional estimator @@ -181,9 +165,10 @@ gradient-equivalent to BPTT outside the assumptions of that approximation. \circ \boldsymbol{\epsilon}^{t'} .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst - ParamDimVjpAlgorithm D_RTRL :class:`D_RTRL` is the concrete, ready-to-use subclass of @@ -216,9 +201,10 @@ it does not allocate multiple rank factors. + (1 - \alpha) \operatorname{diag}(\mathbf{D}_f^t) .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst - IODimVjpAlgorithm pp_prop :class:`pp_prop` is the concrete subclass of :class:`IODimVjpAlgorithm`; @@ -236,23 +222,26 @@ the method is most useful when the recurrent position graph is structurally sparse. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst SnAp -UORO - Random-projection estimator +UORO — Random-projection estimator ---------------------------------- :class:`UORO` carries a rank-one random projection of the full recurrent Jacobian. The projection is an unbiased estimator of the RTRL trace, trading -variance for linear carrier storage. :class:`RandomProjectionVjpAlgorithm` is -the shared engine for algorithms that use this factorization. +variance for linear carrier storage. Its reusable engine, +:class:`RandomProjectionVjpAlgorithm`, is listed under Base Classes. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst - RandomProjectionVjpAlgorithm UORO @@ -266,37 +255,65 @@ gradient to carry credit across finite online windows; :func:`train_synthetic_gradient` updates that predictor. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst ThreeFactor DNI SyntheticGradient .. autosummary:: + :toctree: generated/ :nosignatures: train_synthetic_gradient -SNN Online-Learning Algorithms ------------------------------- +.. _e-prop: + +E-prop — Spiking eligibility-propagation estimator +-------------------------------------------------- -Paper-faithful algorithms tailored to spiking neural networks, all -``ETraceVjpAlgorithm`` subclasses. These are **approximate** (except where a -regime makes them exact); know the regime before relying on their gradients. +:class:`EProp` implements eligibility propagation for recurrent spiking neural +networks, with optional kappa filtering and fixed random-feedback learning +signals. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst EProp + + +.. _ostl: + +OSTL — Recurrent and feedforward estimators +------------------------------------------- + +Online Spatio-Temporal Learning exposes recurrent (with-H) and feedforward +(without-H) regimes as separate concrete classes. + +.. autosummary:: + :toctree: generated/ + :nosignatures: + :template: classtemplate.rst + OSTLRecurrent OSTLFeedforward -Trace helpers reused across the SNN algorithms — a frozen random-feedback -projection and an output-side low-pass filter: + +SNN Helpers +----------- + +Reusable support types for SNN learning signals: a frozen random-feedback +projection and an output-side low-pass filter. .. autosummary:: + :toctree: generated/ :nosignatures: + :template: classtemplate.rst FixedRandomFeedback KappaFilter @@ -334,8 +351,8 @@ Algorithm Comparison - depends on recurrent graph sparsity and ``n`` - Recurrent position graphs whose structural sparsity remains useful over the requested neighbourhood. -Each name above is a thin factory over an :class:`ETraceConfig`; the axes that -distinguish them are: +Each named algorithm above is a preset over an :class:`ETraceConfig`; the axes +that distinguish them are: .. list-table:: :header-rows: 1 diff --git a/docs/apis/compiler.rst b/docs/apis/compiler.rst index 3be5821c..f6ae3da8 100644 --- a/docs/apis/compiler.rst +++ b/docs/apis/compiler.rst @@ -27,9 +27,14 @@ graph object it returns. .. autosummary:: :toctree: generated/ :nosignatures: - :template: classtemplate.rst compile_etrace_graph + +.. autosummary:: + :toctree: generated/ + :nosignatures: + :template: classtemplate.rst + ETraceGraph @@ -57,9 +62,14 @@ Extracts the ``jaxpr`` and state information from a ``brainstate.nn.Module``. .. autosummary:: :toctree: generated/ :nosignatures: - :template: classtemplate.rst extract_module_info + +.. autosummary:: + :toctree: generated/ + :nosignatures: + :template: classtemplate.rst + ModuleInfo @@ -76,6 +86,11 @@ directly from a module. :template: classtemplate.rst HiddenGroup + +.. autosummary:: + :toctree: generated/ + :nosignatures: + find_hidden_groups_from_minfo find_hidden_groups_from_module @@ -95,6 +110,11 @@ ETP primitive is deliberately excluded. :template: classtemplate.rst HiddenParamOpRelation + +.. autosummary:: + :toctree: generated/ + :nosignatures: + find_hidden_param_op_relations_from_minfo find_hidden_param_op_relations_from_module @@ -111,6 +131,11 @@ Perturbation structures used to compute hidden-to-hidden Jacobians :template: classtemplate.rst HiddenPerturbation + +.. autosummary:: + :toctree: generated/ + :nosignatures: + add_hidden_perturbation_from_minfo add_hidden_perturbation_in_module diff --git a/docs/apis/concepts.rst b/docs/apis/concepts.rst index 23773401..b41abccd 100644 --- a/docs/apis/concepts.rst +++ b/docs/apis/concepts.rst @@ -33,7 +33,6 @@ input rank. .. autosummary:: :toctree: generated/ :nosignatures: - :template: classtemplate.rst matmul grouped_matmul diff --git a/docs/conf.py b/docs/conf.py index b9933a99..236f99ad 100644 --- a/docs/conf.py +++ b/docs/conf.py @@ -74,6 +74,17 @@ 'brainx_sphinx_header', ] +# Public API pages document re-exported ``braintrace`` symbols. Keep viewcode +# for objects documented from their defining modules, but do not expose links +# from the public API to private implementation modules. +viewcode_follow_imported_members = False + +# ``brainstate.nn.Module`` uses a metaclass whose generic ``__call__`` would +# otherwise replace concrete constructor signatures with ``(*args, **kwargs)``. +# Render constructors as ``__init__`` methods so autodoc reads the real +# signatures without modifying the runtime classes. +autodoc_class_signature = 'separated' + html_baseurl = 'https://brainx.chaobrain.com/braintrace/' # Keep relative documentation assets local when previewing a built index page. diff --git a/docs/examples/drtrl_examples.rst b/docs/examples/drtrl_examples.rst index 5f9486a0..437c25df 100644 --- a/docs/examples/drtrl_examples.rst +++ b/docs/examples/drtrl_examples.rst @@ -6,19 +6,19 @@ diagonal RTRL implementation across several model and operator families. * `01-basics-integrator.py `__ introduces the integrator workflow. API: - :doc:`D-RTRL `. + :doc:`D-RTRL `. * `02-batching-vmap.py `__ uses per-sample ``vmap`` execution. API: - :doc:`compile `. + :doc:`compile `. * `03-batching-batched.py `__ uses batched primitives. API: - :doc:`compile `. + :doc:`compile `. * `04-vjp-single-step.py `__ demonstrates a single-step VJP. API: - :doc:`D-RTRL `. + :doc:`D-RTRL `. * `05-vjp-multi-step.py `__ demonstrates a multi-step VJP. API: - :doc:`D-RTRL `. + :doc:`D-RTRL `. * `07-operator-lora.py `__ covers a low-rank recurrent operator. API: :doc:`LoRA `. @@ -33,7 +33,7 @@ diagonal RTRL implementation across several model and operator families. :doc:`MiniGRU `. * `11-knob-fast-solve.py `__ examines the ``fast_solve`` implementation option. API: - :doc:`D-RTRL `. + :doc:`D-RTRL `. Read the `D-RTRL examples README `__ @@ -47,6 +47,6 @@ establish general gradient equivalence with BPTT. Related API ----------- -* :doc:`D-RTRL ` +* :doc:`D-RTRL ` * :doc:`ETP Operators ` * :doc:`Algorithm Reference ` diff --git a/docs/examples/pp_prop_examples.rst b/docs/examples/pp_prop_examples.rst index 5238dfd0..52a647ea 100644 --- a/docs/examples/pp_prop_examples.rst +++ b/docs/examples/pp_prop_examples.rst @@ -9,28 +9,28 @@ Follow the numbered scripts in order: * `01-basics-lif-integrator.py `__ introduces the LIF integrator. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `02-neurons-alif-dms.py `__ applies an ALIF model to delayed matching-to-sample. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `03-neurons-gif-working-memory.py `__ demonstrates a GIF working-memory model. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `04-neurons-coba-ei-rsnn.py `__ uses a conductance-based E/I recurrent SNN. API: :doc:`SignedWLinear `. * `05-batching-vmap.py `__ uses per-sample ``vmap`` execution. API: - :doc:`compile `. + :doc:`compile `. * `06-batching-batched.py `__ uses batched primitives. API: - :doc:`compile `. + :doc:`compile `. * `07-vjp-single-step.py `__ demonstrates a single-step VJP. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `08-vjp-multi-step.py `__ demonstrates a multi-step VJP. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `09-operator-sparse.py `__ exercises a sparse recurrent operator. API: :doc:`SparseLinear `. @@ -42,13 +42,13 @@ Follow the numbered scripts in order: :doc:`Conv2d `. * `12-classification-neuromorphic.py `__ provides a neuromorphic classification example. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `13-knob-decay-vs-rank.py `__ compares trace decay with rank. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `14-knob-vjp-method-contrast.py `__ contrasts VJP methods. API: - :doc:`pp-prop `. + :doc:`pp-prop `. Read the `pp-prop examples README `__ @@ -60,6 +60,6 @@ suitability depends on the model dynamics and selected decay or rank. Related API ----------- -* :doc:`pp-prop ` +* :doc:`pp-prop ` * :doc:`ETP Operators ` * :doc:`Algorithm Reference ` diff --git a/docs/examples/rnn_examples.rst b/docs/examples/rnn_examples.rst index fe717cf5..67e65dd4 100644 --- a/docs/examples/rnn_examples.rst +++ b/docs/examples/rnn_examples.rst @@ -22,6 +22,6 @@ workflow for the complete training sequence. Related API ----------- -* :doc:`D-RTRL ` -* :doc:`SnAp ` +* :doc:`D-RTRL ` +* :doc:`SnAp ` * :doc:`Neural Network Layers ` diff --git a/docs/examples/snn_examples.rst b/docs/examples/snn_examples.rst index f34fd2ce..a0900401 100644 --- a/docs/examples/snn_examples.rst +++ b/docs/examples/snn_examples.rst @@ -10,10 +10,10 @@ Learning tasks * `000-lif-snn-for-nmnist.py `__ trains a recurrent LIF-delta SNN on framed N-MNIST events. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `001-gif-snn-for-dms.py `__ trains a GIF recurrent SNN on delayed matching-to-sample. API: - :doc:`pp-prop `. + :doc:`pp-prop `. * `002-coba-ei-rsnn.py `__ trains an excitatory/inhibitory recurrent SNN on an evidence-accumulation task with configurable current- or conductance-based synapses. API: @@ -30,13 +30,13 @@ per-sample ``vmap`` execution: * `003-snn-memory-and-speed-evaluation-all.py `__ runs the complete comparison. API: - :doc:`D-RTRL `. + :doc:`D-RTRL `. * `003-snn-memory-and-speed-evaluation-batched.py `__ uses batched state. API: - :doc:`compile `. + :doc:`compile `. * `003-snn-memory-and-speed-evaluation-vmap.py `__ uses per-sample ``vmap`` execution. API: - :doc:`compile `. + :doc:`compile `. Use these scripts for implementation benchmarking, not as evidence that two algorithms have equivalent gradient accuracy. @@ -47,7 +47,7 @@ workflow before moving to dataset-scale scripts. Related API ----------- -* :doc:`pp-prop ` -* :doc:`EProp ` -* :doc:`OSTLRecurrent ` +* :doc:`pp-prop ` +* :doc:`EProp ` +* :doc:`OSTLRecurrent ` * :doc:`Neural Network Layers `