braintrace 0.2.2__tar.gz → 0.2.3__tar.gz

{braintrace-0.2.2 → braintrace-0.2.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: braintrace
-Version: 0.2.2
+Version: 0.2.3
 Summary: Enabling Scalable Online Learning for Brain Dynamics.
 Author-email: BrainTrace Developers <chao.brain@qq.com>
 License: Apache-2.0 license
@@ -28,6 +28,7 @@ Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: brainstate>=0.5.2
 Requires-Dist: brainunit
+Requires-Dist: brainevent
 Requires-Dist: brainpy.state
 Requires-Dist: braintools
 Provides-Extra: cpu

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/__init__.py

@@ -66,7 +66,9 @@ Examples
 """
 
 
-from typing import TYPE_CHECKING
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Any
 
 from . import nn
 from ._compile import compile
@@ -244,7 +246,7 @@ _DEPRECATED_LEGACY = {
 }
 
 
-def __getattr__(name):
+def __getattr__(name: str) -> Any:
     if name in _DEPRECATED_LEGACY:
         import warnings
         warnings.warn(
@@ -258,5 +260,5 @@ def __getattr__(name):
     raise AttributeError(f'module {__name__!r} has no attribute {name!r}')
 
 
-def __dir__():
+def __dir__() -> list[str]:
     return sorted(list(__all__) + list(_DEPRECATED_LEGACY))

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/__init__.py

@@ -21,6 +21,8 @@ and paper-faithful SNN algorithms (EProp, OSTLRecurrent, OSTLFeedforward, OTPE,
 OTTT, OSTTP).
 """
 
+from __future__ import annotations
+
 from ._common import FixedRandomFeedback, KappaFilter, PresynapticTrace
 from .base import ETraceAlgorithm, EligibilityTrace
 from .d_rtrl import D_RTRL

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/_common.py

@@ -15,6 +15,8 @@
 
 """Shared helpers for the SNN online-learning algorithms."""
 
+from __future__ import annotations
+
 from functools import partial
 from typing import Any, Dict, Optional
 
@@ -23,7 +25,7 @@ import jax
 import jax.numpy as jnp
 import brainunit as u
 
-from braintrace._typing import PyTree
+from braintrace._typing import ArrayLike, PyTree
 
 __all__ = [
     'PresynapticTrace',
@@ -42,12 +44,12 @@ class _ZeroResetState(brainstate.ShortTermState):
     for ``batch_size`` (or prepended for scalar-shaped states).
     """
 
-    def __init__(self, init_value):
+    def __init__(self, init_value: Any) -> None:
         super().__init__(init_value)
         self._init_shape = jnp.shape(init_value)
         self._init_dtype = init_value.dtype
 
-    def reset_state(self, batch_size: Optional[int] = None, **kwargs):
+    def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None:
         """Re-zero the state at its original shape.
 
         Parameters
@@ -105,13 +107,13 @@ class PresynapticTrace(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, leak: float):
+    def __init__(self, init_value: Any, leak: float) -> None:
         super().__init__(init_value)
         if not (0.0 < leak < 1.0):
             raise ValueError(f'leak must be in (0, 1); got {leak}')
         self.leak = float(leak)
 
-    def update(self, x):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one accumulation step :math:`\hat{a} \leftarrow \lambda \cdot \hat{a} + x`.
 
         Parameters
@@ -163,13 +165,13 @@ class KappaFilter(_ZeroResetState):
 
     __module__ = 'braintrace'
 
-    def __init__(self, init_value, kappa: float):
+    def __init__(self, init_value: Any, kappa: float) -> None:
         super().__init__(init_value)
         if not (0.0 <= kappa < 1.0):
             raise ValueError(f'kappa must be in [0, 1); got {kappa}')
         self.kappa = float(kappa)
 
-    def update(self, x):
+    def update(self, x: ArrayLike) -> Any:
         r"""Apply one low-pass step :math:`x_{\mathrm{filt}} \leftarrow (1-\kappa) x + \kappa\, x_{\mathrm{filt}}`.
 
         Parameters
@@ -230,14 +232,14 @@ class FixedRandomFeedback:
 
     __module__ = 'braintrace'
 
-    def __init__(self, n_target: int, n_layer: int, key, init_scale: float = 0.1):
+    def __init__(self, n_target: int, n_layer: int, key: Any, init_scale: float = 0.1) -> None:
         self.B = jax.lax.stop_gradient(
             init_scale * jax.random.normal(key, (n_target, n_layer))
         )
         self.n_target = int(n_target)
         self.n_layer = int(n_layer)
 
-    def project(self, y_target):
+    def project(self, y_target: Any) -> Any:
         """Project the target onto the frozen feedback matrix.
 
         Parameters
@@ -267,7 +269,7 @@ def extract_y_target(args: tuple, *, index: int = -1) -> Optional[jax.Array]:
 def _reset_state_in_a_dict(
     state_dict: Dict[Any, brainstate.State],
     batch_size: Optional[int],
-):
+) -> None:
     """
     Reset the values in a dictionary of states to zero.
 
@@ -293,7 +295,7 @@ def _reset_state_in_a_dict(
 def _zeros_like_batch_or_not(
     batch_size: Optional[int],
     x: jax.Array
-):
+) -> Any:
     """
     Create a zeros array with the same shape and type as the input array,
     optionally including a batch dimension.
@@ -325,7 +327,7 @@ def _batched_zeros_like(
     batch_size: Optional[int],
     num_state: int,  # the number of hidden states
     x: jax.Array  # the input array
-):
+) -> Any:
     """
     Create a batched zeros array with the same shape as the input array,
     extended by the number of hidden states.
@@ -354,7 +356,7 @@ def _batched_zeros_like(
         return u.math.zeros((batch_size, *x.shape, num_state), x.dtype)
 
 
-def _sum_dim(xs: jax.Array, axis: int = -1):
+def _sum_dim(xs: jax.Array, axis: int = -1) -> Any:
     """
     Sums the elements along the last dimension of each array in a PyTree.
 
@@ -373,7 +375,7 @@ def _sum_dim(xs: jax.Array, axis: int = -1):
     return jax.tree.map(lambda x: u.math.sum(x, axis=axis), xs)
 
 
-def _unit_safe_add(a, b):
+def _unit_safe_add(a: Any, b: Any) -> Any:
     """Add two leaves, stripping units only when one side has units and the other does not.
 
     Gradient contributions for the same weight may come from paths that
@@ -390,7 +392,7 @@ def _unit_safe_add(a, b):
     return u.math.add(a, b)
 
 
-def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
+def _extract_leaf(pytree_val: PyTree, leaf_idx: int) -> Any:
     """Return the leaf at ``leaf_idx`` in ``jax.tree.leaves(pytree_val)``.
 
     Bare arrays (treedef with a single leaf) return the array unchanged.
@@ -409,7 +411,7 @@ def _extract_leaf(pytree_val: PyTree, leaf_idx: int):
 def _wrap_leaves_as_pytree(
     reference_pytree: PyTree,
     leaf_grads: Dict[int, jax.Array],
-):
+) -> Any:
     """Build a pytree matching ``reference_pytree`` with ``leaf_grads``
     inserted at the given leaf indices; any other leaf is zero-filled.
 
@@ -441,7 +443,7 @@ def _wrap_leaves_as_pytree(
 
 
 def _route_grads_by_path(
-    relation,
+    relation: Any,
     per_key_grads: Dict[str, jax.Array],
     weight_vals: Dict[Any, PyTree],
     target_dict: Dict[Any, PyTree],
@@ -478,7 +480,7 @@ def _update_dict(
     key: Any,
     value: PyTree,
     error_when_no_key: Optional[bool] = False
-):
+) -> None:
     """Update the dictionary.
 
     If the key exists, then add the value to the existing value.

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/base.py

@@ -19,6 +19,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Dict, Any, Optional
 
 import brainstate
@@ -89,7 +91,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         model: brainstate.nn.Module,
         graph_executor: ETraceGraphExecutor,
         name: Optional[str] = None,
-    ):
+    ) -> None:
         super().__init__(name=name)  # type: ignore[call-arg]  # brainstate hides Module.__init__ from type checkers
 
         # the model
@@ -210,7 +212,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         assert self._other_states is not None
         return self._other_states
 
-    def _split_state(self):
+    def _split_state(self) -> None:
         # --- the state separation --- #
         #
         # [NOTE]
@@ -228,7 +230,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
             self._other_states
         ) = self.graph.module_info.retrieved_model_states.split(brainstate.ParamState, brainstate.HiddenState, ...)
 
-    def compile_graph(self, *args) -> None:
+    def compile_graph(self, *args: Any) -> None:
         r"""
         Compile the eligibility trace graph of the relationship between etrace weights, states and operators.
 
@@ -318,7 +320,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.graph_executor.show_graph(verbose=verbose, return_msg=return_msg)
 
-    def __call__(self, *args) -> Any:
+    def __call__(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -334,7 +336,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         return self.update(*args)
 
-    def update(self, *args) -> Any:
+    def update(self, *args: Any) -> Any:
         """
         Update the model and the eligibility trace states.
 
@@ -355,7 +357,7 @@ class ETraceAlgorithm(brainstate.nn.Module):
         """
         raise NotImplementedError
 
-    def init_etrace_state(self, *args, **kwargs) -> None:
+    def init_etrace_state(self, *args: Any, **kwargs: Any) -> None:
         """
         Initialize the eligibility trace states of the etrace algorithm.
 

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/d_rtrl.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 # ==============================================================================
 
+from __future__ import annotations
+
 from .param_dim_vjp import ParamDimVjpAlgorithm
 
 __all__ = [

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/d_rtrl_test.py

@@ -25,7 +25,6 @@ from braintrace._algorithm.d_rtrl import D_RTRL
 from braintrace._algorithm.param_dim_vjp import (
     ParamDimVjpAlgorithm,
     _remove_units,
-    _fast_solve_contract,
 )
 from braintrace._etrace_model_test import (
     IF_Delta_Dense_Layer,
@@ -39,7 +38,7 @@ from braintrace._etrace_model_test import (
     ALIF_STDExpCu_Dense_Layer,
     ALIF_STPExpCu_Dense_Layer,
 )
-from braintrace._op import etp_mm_p
+from braintrace._op import etp_mm_p, get_fast_path_rules
 
 
 # ---------------------------------------------------------------------------
@@ -657,7 +656,7 @@ class TestDRtrlDictTraceStorage:
 # Helpers for the trace-update specialization + dtype-knob tests
 # ---------------------------------------------------------------------------
 
-def _rnn_mm(seed=0, n=4, bias=False):
+def _rnn_mm(seed=0, n=4, bias=False, weight_fn=None, bias_fn=None):
     """Batched single-recurrent-weight tanh RNN -> etp_mm_p, num_state==1."""
 
     class Net(brainstate.nn.Module):
@@ -670,7 +669,10 @@ def _rnn_mm(seed=0, n=4, bias=False):
 
         def update(self, x):
             b = self.b.value if self.b is not None else None
-            self.h.value = jax.nn.tanh(braintrace.matmul(self.h.value + x, self.w.value, b))
+            self.h.value = jax.nn.tanh(braintrace.matmul(
+                self.h.value + x, self.w.value, b,
+                weight_fn=weight_fn, bias_fn=bias_fn,
+            ))
             return self.h.value
 
     net = Net()
@@ -678,7 +680,7 @@ def _rnn_mm(seed=0, n=4, bias=False):
     return net
 
 
-def _rnn_mv(seed=0, n=4):
+def _rnn_mv(seed=0, n=4, weight_fn=None):
     """Unbatched single-recurrent-weight tanh RNN -> etp_mv_p, num_state==1."""
 
     class Net(brainstate.nn.Module):
@@ -690,7 +692,9 @@ def _rnn_mv(seed=0, n=4):
             self.h = brainstate.HiddenState(jnp.zeros((n,)))
 
         def update(self, x):
-            self.h.value = jax.nn.tanh(braintrace.matmul(self.h.value + x, self.w.value))
+            self.h.value = jax.nn.tanh(braintrace.matmul(
+                self.h.value + x, self.w.value, weight_fn=weight_fn,
+            ))
             return self.h.value
 
     net = Net()
@@ -698,6 +702,27 @@ def _rnn_mv(seed=0, n=4):
     return net
 
 
+def _rnn_elemwise(seed=0, n=4, weight_fn=None):
+    """Single element-wise-weight tanh RNN -> etp_elemwise_p, num_state==1."""
+
+    class Net(brainstate.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.alpha = brainstate.ParamState(
+                0.5 * jax.random.normal(jax.random.PRNGKey(seed), (n,))
+            )
+            self.h = brainstate.HiddenState(jnp.zeros((1, n)))
+
+        def update(self, x):
+            a = braintrace.element_wise(self.alpha.value, weight_fn=weight_fn)
+            self.h.value = jax.nn.tanh(a * (self.h.value + x))
+            return self.h.value
+
+    net = Net()
+    brainstate.nn.init_all_states(net, batch_size=1)
+    return net
+
+
 def _run_updates(algo, xs):
     """Compile, init, run len(xs) eager updates; return etrace_bwg snapshot.
 
@@ -773,6 +798,73 @@ class TestS1RecurrentBroadcastEqualsLegacy:
         _assert_etrace_lists_equal(fast, legacy, exact=True)
 
 
+# ---------------------------------------------------------------------------
+# weight_fn/bias_fn — the fast closed-form trace must equal the rule (legacy)
+# path when a transform hook is present.
+#
+# The fast kernels compute the instant term as the outer product x ⊗ df, which
+# equals ∂h/∂V (the *transformed* weight). With a transform present the true
+# eligibility trace is ∂h/∂W_raw = (x ⊗ df) ⊙ f'(W); the f' factor is supplied
+# only by the rule path's ``xy_to_dw`` (via jax.vjp). So the fast path must fall
+# back to the rules whenever a ``*_fn`` hook is set, otherwise the trace is the
+# gradient w.r.t. the transformed weight and silently drops f'(W).
+# ---------------------------------------------------------------------------
+
+class TestTransformFastEqualsLegacy:
+    """With a weight_fn/bias_fn hook, the default (fast) trace must match the
+    rule path bit-for-bit — i.e. the fast path falls back to the rules so the
+    f'(W) factor from ``xy_to_dw`` is not dropped.
+
+    Run several steps so the recurrent term (non-zero from the 2nd update) is
+    exercised, not just the instant term."""
+
+    def _xs_mm(self, n=4, steps=4):
+        return brainstate.random.randn(steps, 1, n)
+
+    def _xs_mv(self, n=4, steps=4):
+        return brainstate.random.randn(steps, n)
+
+    def test_mm_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        wfn = lambda w: w ** 2
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(weight_fn=wfn), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(weight_fn=wfn), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_mm_weight_fn_and_bias_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        fast = _run_updates(
+            ParamDimVjpAlgorithm(
+                _rnn_mm(bias=True, weight_fn=jnp.tanh, bias_fn=lambda b: b * 2.0),
+                fast_solve=True), xs)
+        legacy = _run_updates(
+            ParamDimVjpAlgorithm(
+                _rnn_mm(bias=True, weight_fn=jnp.tanh, bias_fn=lambda b: b * 2.0),
+                fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_mv_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mv()
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mv(weight_fn=jnp.tanh), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mv(weight_fn=jnp.tanh), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_elemwise_weight_fn_fast_equals_legacy(self):
+        xs = self._xs_mm()
+        wfn = jnp.tanh
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_elemwise(weight_fn=wfn), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_elemwise(weight_fn=wfn), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+    def test_no_transform_still_uses_fast_path(self):
+        """Guard: without a transform, fast and legacy remain bit-identical (the
+        gate must not disable the fast path when no ``*_fn`` hook is present)."""
+        xs = self._xs_mm()
+        fast = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(), fast_solve=True), xs)
+        legacy = _run_updates(ParamDimVjpAlgorithm(_rnn_mm(), fast_solve=False), xs)
+        _assert_etrace_lists_equal(fast, legacy, exact=True)
+
+
 # ---------------------------------------------------------------------------
 # Change A — trace_dtype knob (store the eligibility trace in lower precision)
 # ---------------------------------------------------------------------------
@@ -838,7 +930,8 @@ class TestTraceDtypeKnob:
 
 def _alif(n_in=4, n_rec=5):
     """Adaptive-LIF dense layer -> coupled (V, adaptation) hidden group with
-    num_state==2, exercising the S>1 einsum branch of _fast_recurrent_term."""
+    num_state==2, exercising the S>1 einsum branch of the dense fast-path
+    recurrent rule (``FastPathRules.recurrent``)."""
     model = ALIF_Delta_Dense_Layer(n_in, n_rec)
     brainstate.nn.init_all_states(model)
     return model
@@ -891,11 +984,12 @@ class TestSolveBatchFold:
             'bias': brainstate.random.randn(B, O, S),
         }
         # reference: current per-batch contraction, then explicit batch sum
-        ref = _fast_solve_contract(etp_mm_p, diag_like, etrace)
+        solve = get_fast_path_rules(etp_mm_p).solve
+        ref = solve(diag_like, etrace, fold_batch=False)
         ref_w = ref['weight'].sum(axis=0)   # (I, O)
         ref_b = ref['bias'].sum(axis=0)     # (O,)
         # new: fold the batch axis inside the einsum
-        folded = _fast_solve_contract(etp_mm_p, diag_like, etrace, fold_batch=True)
+        folded = solve(diag_like, etrace, fold_batch=True)
         assert folded['weight'].shape == (I, O)
         assert folded['bias'].shape == (O,)
         npt.assert_allclose(folded['weight'], ref_w, atol=1e-6)

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/e_prop.py

@@ -29,7 +29,9 @@ that make the rule biologically plausible:
 See :class:`EProp` for the mathematical formulation, references, and an example.
 """
 
-from typing import Dict, Optional
+from __future__ import annotations
+
+from typing import Any, Dict, Optional
 
 import brainstate
 import jax
@@ -158,8 +160,8 @@ class EProp(ParamDimVjpAlgorithm):
         name: Optional[str] = None,
         vjp_method: str = 'single-step',
         fast_solve: bool = True,
-        **kwargs,
-    ):
+        **kwargs: Any,
+    ) -> None:
         if feedback not in ('symmetric', 'random'):
             raise ValueError(
                 f"feedback must be 'symmetric' or 'random'; got {feedback!r}"
@@ -181,7 +183,7 @@ class EProp(ParamDimVjpAlgorithm):
         self._kappa_filters: Dict[int, KappaFilter] = {}
         self._random_feedback: Dict[int, FixedRandomFeedback] = {}
 
-    def init_etrace_state(self, *args, **kwargs):
+    def init_etrace_state(self, *args: Any, **kwargs: Any) -> None:
         super().init_etrace_state(*args, **kwargs)
         self._kappa_filters = {}
         self._random_feedback = {}
@@ -197,6 +199,7 @@ class EProp(ParamDimVjpAlgorithm):
                     )
         if self.feedback == 'random':
             key = self._random_feedback_key
+            assert key is not None  # constructor enforces a key when feedback='random'
             for rel in self.graph.hidden_param_op_relations:
                 for group in rel.hidden_groups:
                     gid = group.index
@@ -209,17 +212,17 @@ class EProp(ParamDimVjpAlgorithm):
                         n_target=n_layer, n_layer=n_layer, key=sub, init_scale=0.1
                     )
 
-    def reset_state(self, batch_size: Optional[int] = None, **kwargs):
+    def reset_state(self, batch_size: Optional[int] = None, **kwargs: Any) -> None:
         super().reset_state(batch_size=batch_size, **kwargs)
         for flt in self._kappa_filters.values():
             flt.reset_state(batch_size=batch_size)
 
-    def _compute_learning_signal(self, dl_autodiff, args):
+    def _compute_learning_signal(self, dl_autodiff: Any, args: Any) -> Any:
         signals = list(dl_autodiff)
         if self.feedback == 'random' and self._random_feedback:
             # dl_autodiff[g].shape == (*varshape, num_state). Project over the
             # trailing n_layer axis (-2), preserving num_state on the tail.
-            def _project(B, s):
+            def _project(B: Any, s: Any) -> Any:
                 return jnp.einsum('...lj,lk->...kj', s, B)
 
             signals = [
@@ -231,7 +234,7 @@ class EProp(ParamDimVjpAlgorithm):
             # KappaFilter state carries varshape, but signal has an extra
             # trailing num_state axis. Collapse num_state for filter purposes;
             # broadcast the filtered value back.
-            def _filter(flt, s):
+            def _filter(flt: Any, s: Any) -> Any:
                 # collapse num_state tail: sum over last axis produces shape (*varshape,)
                 collapsed = s.sum(axis=-1)
                 filtered = flt.update(collapsed)

{braintrace-0.2.2 → braintrace-0.2.3}/braintrace/_algorithm/graph_executor.py

@@ -35,6 +35,8 @@
 
 # -*- coding: utf-8 -*-
 
+from __future__ import annotations
+
 from typing import Dict, Any, Optional
 
 import brainstate
@@ -70,7 +72,7 @@ class ETraceGraphExecutor:
         self,
         model: brainstate.nn.Module,
         include_recurrent_mixing: bool = False,
-    ):
+    ) -> None:
         # The original model
         if not isinstance(model, brainstate.nn.Module):
             raise TypeError(
@@ -152,7 +154,7 @@ class ETraceGraphExecutor:
             self._state_id_to_path = {id(state): path for path, state in self.states.items()}
         return self._state_id_to_path
 
-    def compile_graph(self, *args) -> None:
+    def compile_graph(self, *args: Any) -> None:
         r"""
         Build the eligibility trace graph for the model based on the provided inputs.
 
@@ -217,7 +219,7 @@ class ETraceGraphExecutor:
 
     def solve_h2w_h2h_jacobian(
         self,
-        *args,
+        *args: Any,
     ) -> Any:
         r"""
         Compute the hidden-to-weight and hidden-to-hidden Jacobian matrices.
@@ -262,7 +264,7 @@ class ETraceGraphExecutor:
                                   'implemented in the subclass.')
 
     def solve_h2w_h2h_l2h_jacobian(
-        self, *args,
+        self, *args: Any,
     ) -> Any:
         r"""
         Compute the hidden-to-weight and hidden-to-hidden Jacobian matrices, along with the VJP transformed