Add the exact same-step gradient for outputs that bypass new hidden values (#185)

braintrace/_algorithm/io_dim_vjp.py

@@ -1106,6 +1106,35 @@ def _update_etrace_data(
 
         return hist_etrace_vals
 
+    def _solve_previous_trace_gradients(
+        self,
+        running_index: int,
+        etrace_h2w_at_t_minus_1: Tuple[
+            Dict[ETraceX_Key, jax.Array],
+            Dict[ETraceDF_Key, jax.Array]
+        ],
+        dl_to_hidden_groups: Sequence[jax.Array],
+        weight_vals: Dict[Path, PyTree],
+    ) -> Dict[Path, PyTree]:
+        """Contract the previous trace with the hidden output-bypass signal.
+
+        See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`. The
+        previous trace is one update younger, so its warm-up bias correction
+        is indexed by ``running_index - 1``.
+        """
+        dG_weights: Dict[Path, Any] = {path: None for path in self.param_states.keys()}
+        _solve_IO_dim_weight_gradients(
+            etrace_h2w_at_t_minus_1,
+            dG_weights,
+            dl_to_hidden_groups,
+            self.graph.hidden_param_op_relations,
+            weight_vals,
+            running_index - 1,
+            self.decay_f,
+            fast_solve=self.fast_solve,
+        )
+        return dG_weights
+
     def _solve_weight_gradients(
         self,
         running_index: int,

braintrace/_algorithm/param_dim_vjp.py

@@ -1408,6 +1408,35 @@ def _update_etrace_data(
 
         return new_etrace
 
+    def _solve_previous_trace_gradients(
+        self,
+        running_index: int,
+        etrace_h2w_at_t_minus_1: Dict[ETraceWG_Key, PyTree],
+        dl_to_hidden_groups: Sequence[jax.Array],
+        weight_vals: Dict[Path, PyTree],
+    ) -> Dict[Path, PyTree]:
+        """Contract the previous trace with the hidden output-bypass signal.
+
+        See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`. With
+        trace filters active, the previous *filtered* trace (the filter's
+        current value, not yet advanced this step) is the one paired.
+        """
+        etrace = dict(etrace_h2w_at_t_minus_1)
+        if self._trace_filters:
+            for key, flt in self._trace_filters.items():
+                if key in etrace:
+                    etrace[key] = flt.value
+        dG_weights: Dict[Path, Any] = {path: None for path in self.param_states}
+        _solve_param_dim_weight_gradients(
+            etrace,
+            dG_weights,
+            dl_to_hidden_groups,
+            self.graph.hidden_param_op_relations,
+            weight_vals,
+            fast_solve=self.fast_solve,
+        )
+        return dG_weights
+
     def _solve_weight_gradients(
         self,
         running_index: int,

braintrace/_algorithm/random_projection_vjp.py

@@ -665,6 +665,36 @@ def _update_etrace_data(
     # the boundary contraction
     # ------------------------------------------------------------------ #
 
+    def _solve_previous_trace_gradients(
+        self,
+        running_index: int,
+        etrace_h2w_at_t_minus_1: Dict[str, Any],
+        dl_to_hidden_groups: Sequence[jax.Array],
+        weight_vals: Dict[Path, PyTree],
+    ) -> Dict[Path, PyTree]:
+        """Contract the previous trace with the hidden output-bypass signal.
+
+        See :meth:`ETraceVjpAlgorithm._solve_previous_trace_gradients`.
+        """
+        return self._contract_trace(etrace_h2w_at_t_minus_1, dl_to_hidden_groups)
+
+    def _contract_trace(
+        self,
+        etrace: Dict[str, Any],
+        dl_to_hidden_groups: Sequence[jax.Array],
+    ) -> Dict[Path, Any]:
+        # ``sum_g (signal_g . s_tilde_g) * theta_tilde_g``
+        dG_weights: Dict[Path, Any] = {path: None for path in self.param_states}
+        for group in self.graph.hidden_groups:
+            gi = group.index
+            signal = u.get_mantissa(dl_to_hidden_groups[gi])
+            s_tilde = etrace[_S_TILDE][gi]
+            scale = jnp.sum(signal * s_tilde)
+            for path in self._etp_paths_of_group(group):
+                theta = etrace[_THETA_TILDE][(gi, path)]
+                _update_dict(dG_weights, path, _scale_tree(theta, scale))
+        return dG_weights
+
     def _solve_weight_gradients(
         self,
         running_index: int,
@@ -679,22 +709,14 @@ def _solve_weight_gradients(
         A scalar per group times a parameter-shaped array — the whole reason the
         parameter factor carries no hidden axis.
         """
-        dG_weights: Dict[Path, Any] = {path: None for path in self.param_states}
-
-        for group in self.graph.hidden_groups:
-            gi = group.index
-            signal = u.get_mantissa(dl_to_hidden_groups[gi])
-            s_tilde = etrace_h2w_at_t[_S_TILDE][gi]
-            scale = jnp.sum(signal * s_tilde)
-            for path in self._etp_paths_of_group(group):
-                theta = etrace_h2w_at_t[_THETA_TILDE][(gi, path)]
-                _update_dict(dG_weights, path, _scale_tree(theta, scale))
+        dG_weights = self._contract_trace(etrace_h2w_at_t, dl_to_hidden_groups)
 
         # the non-etrace weight gradients (reverse-AD, in-window)
         for path, dg in dl_to_nonetws_at_t.items():
             _update_dict(dG_weights, path, dg)
 
-        # the in-window direct term for etrace parameters (multi-step only)
+        # the in-window direct term for etrace parameters (multi-step), or
+        # the same-step output-bypass term (single-step, #185)
         if dl_to_etws_at_t is not None:
             for path, dg in dl_to_etws_at_t.items():
                 _update_dict(dG_weights, path, dg, error_when_no_key=True)