From aad33612fcda8dbde6bd07f5726933565dad8e4a Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Fri, 24 Jul 2026 09:34:19 +0900
Subject: [PATCH 1/8] Reduce CPU->GPU transfer every time step with learning.

---
 bindsnet/learning/MCC_learning.py | 29 ++++++++++++++++++++++-------
 1 file changed, 22 insertions(+), 7 deletions(-)

diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py
index 33ebacd0a..3fe2014c9 100644
--- a/bindsnet/learning/MCC_learning.py
+++ b/bindsnet/learning/MCC_learning.py
@@ -503,13 +503,28 @@ def _connection_update(self, **kwargs) -> None:
         source_s = self.source.s.view(batch_size, -1).float()
         target_s = self.target.s.view(batch_size, -1).float()
 
-        # Parse keyword arguments.
+        # Reward from current time step
         reward = kwargs["reward"]
-        a_plus = torch.tensor(
-            kwargs.get("a_plus", 1.0), device=self.feature_value.device
+
+        # Build learning-rate and decay tensors
+        if not hasattr(self, "_a_plus_default"):
+            dev = self.feature_value.device
+            self._a_plus_default = torch.tensor(1.0, device=dev)
+            self._a_minus_default = torch.tensor(-1.0, device=dev)
+            dt = float(self.connection.dt)
+            self._decay_plus = torch.exp(-dt / self.tc_plus.to(dev))
+            self._decay_minus = torch.exp(-dt / self.tc_minus.to(dev))
+        a_plus = kwargs.get("a_plus", None)
+        a_plus = (
+            self._a_plus_default
+            if a_plus is None
+            else torch.as_tensor(a_plus, device=self.feature_value.device)
         )
-        a_minus = torch.tensor(
-            kwargs.get("a_minus", -1.0), device=self.feature_value.device
+        a_minus = kwargs.get("a_minus", None)
+        a_minus = (
+            self._a_minus_default
+            if a_minus is None
+            else torch.as_tensor(a_minus, device=self.feature_value.device)
         )
 
         # Compute weight update based on the eligibility value of the past timestep.
@@ -535,9 +550,9 @@ def _connection_update(self, **kwargs) -> None:
             self.feature_value += update
 
         # Update P^+ and P^- values.
-        self.p_plus *= torch.exp(-self.connection.dt / self.tc_plus)
+        self.p_plus *= self._decay_plus
         self.p_plus += a_plus * source_s
-        self.p_minus *= torch.exp(-self.connection.dt / self.tc_minus)
+        self.p_minus *= self._decay_minus
         self.p_minus += a_minus * target_s
 
         # Calculate point eligibility value.

From 1b6e4a2e56d46106b7f4b18c32adf6fd52c8b435 Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Fri, 24 Jul 2026 10:00:06 +0900
Subject: [PATCH 2/8] Foldable pipelines

---
 bindsnet/network/topology.py | 62 ++++++++++++++++++++++++++++++++++++
 1 file changed, 62 insertions(+)

diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py
index 10df45824..42d687030 100644
--- a/bindsnet/network/topology.py
+++ b/bindsnet/network/topology.py
@@ -413,6 +413,7 @@ def __init__(
         pipeline: list = [],
         manual_update: bool = False,
         traces: bool = False,
+        compute_dtype: Optional[torch.dtype] = None,
         **kwargs,
     ) -> None:
         # language=rst
@@ -426,6 +427,8 @@ def __init__(
         :param manual_update: Set to :code:`True` to disable automatic updates (applying learning rules) to connection features.
             False by default, updates called after each time step
         :param traces: Set to :code:`True` to record history of connection activity (for monitors)
+        :param compute_dtype: Optional low-precision dtype (e.g. ``torch.bfloat16``)
+            for the matmul fast path only; ``None`` (default) keeps full ``float32``.
         """
 
         super().__init__(source, target, device, pipeline, **kwargs)
@@ -434,6 +437,49 @@ def __init__(
         if self.traces:
             self.activity = None
 
+        self._w_eff = None          # 'Folded' weight matrix for faster computing
+        self._has_learning = None
+        self.compute_dtype = compute_dtype
+
+    def _pipeline_has_learning(self) -> bool:
+        # language=rst
+        """Whether any pipeline feature carries a real (non-``NoOp``) learning
+        rule, i.e. whether the folded-weight cache must be invalidated on update."""
+        if self._has_learning is None:
+            from ..learning.MCC_learning import NoOp
+
+            self._has_learning = any(
+                not isinstance(getattr(f, "learning_rule", None), NoOp)
+                for f in self.pipeline
+            )
+        return self._has_learning
+
+    def _folded_weight(self) -> Optional[torch.Tensor]:
+        # language=rst
+        """
+        Return the product of the pipeline's foldable feature values. If the
+        pipeline is empty or contains a non-foldable feature, returns none.
+        """
+        if not self.pipeline:
+            return None
+        if self._w_eff is not None:
+            return self._w_eff
+        w_eff = None
+        for f in self.pipeline:
+            v = f.matmul_fold_value()
+            if v is None:
+                return None  # a non-foldable feature -> use the generic path
+            w_eff = v if w_eff is None else w_eff * v
+        # A product of only boolean masks would be non-float; force the matmul
+        # dtype (float32 by default, or the opt-in low-precision compute_dtype).
+        if self.compute_dtype is not None:
+            w_eff = w_eff.to(self.compute_dtype)
+        elif not torch.is_floating_point(w_eff):
+            w_eff = w_eff.float()
+        if not self.manual_update:
+            self._w_eff = w_eff
+        return w_eff
+
     def compute(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
         """
@@ -444,6 +490,18 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
                  decaying spike activation).
         """
 
+        # Fast path: when the whole pipeline folds to a static elementwise-multiply
+        w_eff = self._folded_weight()
+        if w_eff is not None:
+            out_signal = s.view(s.size(0), self.source.n).to(w_eff.dtype) @ w_eff
+            if out_signal.dtype != torch.float32:
+                out_signal = out_signal.float()
+            if self.traces:
+                self.activity = out_signal
+            if out_signal.size() != torch.Size([s.size(0)] + self.target.shape):
+                return out_signal.view(s.size(0), *self.target.shape)
+            return out_signal
+
         # Change to numeric type (torch doesn't like booleans for matrix ops)
         # Note: .float() is an expensive operation. Use as minimally as possible!
         # if s.dtype != torch.float32:
@@ -516,6 +574,8 @@ def update(self, **kwargs) -> None:
             # Pipeline learning
             for f in self.pipeline:
                 f.update(**kwargs)
+            if self._pipeline_has_learning():
+                self._w_eff = None  # weights changed -> rebuild fold next compute
 
     def normalize(self) -> None:
         # language=rst
@@ -525,6 +585,7 @@ def normalize(self) -> None:
         # Normalize pipeline features
         for f in self.pipeline:
             f.normalize()
+        self._w_eff = None  # normalization may change weights -> invalidate fold
 
     def reset_state_variables(self) -> None:
         # language=rst
@@ -535,6 +596,7 @@ def reset_state_variables(self) -> None:
 
         for f in self.pipeline:
             f.reset_state_variables()
+        self._w_eff = None  # rebuild the fold cache at the next sample
 
 
 class Conv1dConnection(AbstractConnection):

From 2428476a5d8d4aea4561bd70cabfad307e04c68f Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Fri, 24 Jul 2026 10:01:02 +0900
Subject: [PATCH 3/8] More files for foldable pipelines

---
 bindsnet/network/topology_features.py | 29 +++++++++++++++++++++++++++
 1 file changed, 29 insertions(+)

diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py
index 43ccc1389..c39288020 100644
--- a/bindsnet/network/topology_features.py
+++ b/bindsnet/network/topology_features.py
@@ -170,6 +170,21 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
         """
         pass
 
+    def matmul_fold_value(self) -> Optional[torch.Tensor]:
+        # language=rst
+        """
+        Fast-path hook for :class:`MulticompartmentConnection`.
+
+        If this feature is a *pure elementwise multiply* by a static
+        ``[source.n, target.n]`` tensor -- i.e. ``compute(x) == x * V`` with no
+        per-call side effects on ``V`` -- return that ``V`` so the connection can
+        fold the whole pipeline into one weight matrix and evaluate as a matmul
+        (``out = s @ W_eff``) instead of materialising ``[B, source.n,
+        target.n]``. Returning ``None`` (the safe default) forces the generic
+        pipeline path.
+        """
+        return None
+
     def prime_feature(self, connection, device, **kwargs) -> None:
         # language=rst
         """
@@ -507,6 +522,12 @@ def __init__(
     def compute(self, conn_spikes) -> torch.Tensor:
         return conn_spikes * self.value
 
+    def matmul_fold_value(self) -> Optional[torch.Tensor]:
+        # A boolean mask is a pure elementwise multiply (True->1, False->0).
+        if getattr(self, "sparse", False):
+            return None
+        return self.value
+
     def reset_state_variables(self) -> None:
         pass
 
@@ -644,6 +665,14 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
 
         return return_val
 
+    def matmul_fold_value(self) -> Optional[torch.Tensor]:
+        # Foldable only as a plain multiply: skip when polarity enforcement or
+        # per-time-step normalisation would mutate ``value`` during compute, or
+        # when the value is stored sparse.
+        if self.enforce_polarity or self.norm_frequency == "time step" or self.sparse:
+            return None
+        return self.value
+
     def prime_feature(self, connection, device, **kwargs) -> None:
         #### Initialize value ####
         if self.value is None:

From 677d196e47895a47d60827e5fbc24b2bc8f5c62c Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Sat, 25 Jul 2026 03:38:40 +0900
Subject: [PATCH 4/8] Sparse compute

---
 bindsnet/network/topology.py | 24 +++++++++++++++++++++++-
 1 file changed, 23 insertions(+), 1 deletion(-)

diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py
index 42d687030..07ee8a5b6 100644
--- a/bindsnet/network/topology.py
+++ b/bindsnet/network/topology.py
@@ -414,6 +414,7 @@ def __init__(
         manual_update: bool = False,
         traces: bool = False,
         compute_dtype: Optional[torch.dtype] = None,
+        sparse_compute: bool = False,
         **kwargs,
     ) -> None:
         # language=rst
@@ -429,6 +430,8 @@ def __init__(
         :param traces: Set to :code:`True` to record history of connection activity (for monitors)
         :param compute_dtype: Optional low-precision dtype (e.g. ``torch.bfloat16``)
             for the matmul fast path only; ``None`` (default) keeps full ``float32``.
+        :param sparse_compute: Set to :code:`True` to compute sparse activity (<~20% neurons active)
+        efficiently
         """
 
         super().__init__(source, target, device, pipeline, **kwargs)
@@ -440,6 +443,7 @@ def __init__(
         self._w_eff = None          # 'Folded' weight matrix for faster computing
         self._has_learning = None
         self.compute_dtype = compute_dtype
+        self.sparse_compute = sparse_compute
 
     def _pipeline_has_learning(self) -> bool:
         # language=rst
@@ -480,6 +484,20 @@ def _folded_weight(self) -> Optional[torch.Tensor]:
             self._w_eff = w_eff
         return w_eff
 
+    def _sparse_matmul(self, s_flat: torch.Tensor, w_eff: torch.Tensor) -> torch.Tensor:
+        # language=rst
+        """
+        Activity-sparse form of ``s_flat @ w_eff``. Read only
+        the rows of ``w_eff`` for source neurons that spiked this step. Numerically
+        identical to the dense matmul. More efficient when few source neurons are active.
+        """
+        active = s_flat.any(dim=0).nonzero(as_tuple=False).squeeze(1)
+        if active.numel() == 0:
+            return torch.zeros(
+                s_flat.size(0), self.target.n, device=w_eff.device, dtype=w_eff.dtype
+            )
+        return s_flat[:, active].to(w_eff.dtype) @ w_eff.index_select(0, active)
+
     def compute(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
         """
@@ -493,7 +511,11 @@ def compute(self, s: torch.Tensor) -> torch.Tensor:
         # Fast path: when the whole pipeline folds to a static elementwise-multiply
         w_eff = self._folded_weight()
         if w_eff is not None:
-            out_signal = s.view(s.size(0), self.source.n).to(w_eff.dtype) @ w_eff
+            s_flat = s.view(s.size(0), self.source.n)
+            if self.sparse_compute:
+                out_signal = self._sparse_matmul(s_flat, w_eff)
+            else:
+                out_signal = s_flat.to(w_eff.dtype) @ w_eff
             if out_signal.dtype != torch.float32:
                 out_signal = out_signal.float()
             if self.traces:

From 884010b09253cfe322c25ebd6ba1d155314fafb3 Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Sun, 26 Jul 2026 16:07:14 +0900
Subject: [PATCH 5/8] Folding pipeline computation

---
 bindsnet/network/topology.py          | 173 +++++++++-----------------
 bindsnet/network/topology_features.py | 113 +++++++++--------
 2 files changed, 115 insertions(+), 171 deletions(-)

diff --git a/bindsnet/network/topology.py b/bindsnet/network/topology.py
index 07ee8a5b6..eaa4485f8 100644
--- a/bindsnet/network/topology.py
+++ b/bindsnet/network/topology.py
@@ -413,7 +413,6 @@ def __init__(
         pipeline: list = [],
         manual_update: bool = False,
         traces: bool = False,
-        compute_dtype: Optional[torch.dtype] = None,
         sparse_compute: bool = False,
         **kwargs,
     ) -> None:
@@ -428,10 +427,8 @@ def __init__(
         :param manual_update: Set to :code:`True` to disable automatic updates (applying learning rules) to connection features.
             False by default, updates called after each time step
         :param traces: Set to :code:`True` to record history of connection activity (for monitors)
-        :param compute_dtype: Optional low-precision dtype (e.g. ``torch.bfloat16``)
-            for the matmul fast path only; ``None`` (default) keeps full ``float32``.
-        :param sparse_compute: Set to :code:`True` to compute sparse activity (<~20% neurons active)
-        efficiently
+        :param sparse_compute: Set to :code:`True` to read only the rows of the effective
+            weight for source neurons that spiked (a win only when few are active).
         """
 
         super().__init__(source, target, device, pipeline, **kwargs)
@@ -440,123 +437,75 @@ def __init__(
         if self.traces:
             self.activity = None
 
-        self._w_eff = None          # 'Folded' weight matrix for faster computing
-        self._has_learning = None
-        self.compute_dtype = compute_dtype
         self.sparse_compute = sparse_compute
 
-    def _pipeline_has_learning(self) -> bool:
-        # language=rst
-        """Whether any pipeline feature carries a real (non-``NoOp``) learning
-        rule, i.e. whether the folded-weight cache must be invalidated on update."""
-        if self._has_learning is None:
-            from ..learning.MCC_learning import NoOp
-
-            self._has_learning = any(
-                not isinstance(getattr(f, "learning_rule", None), NoOp)
-                for f in self.pipeline
-            )
-        return self._has_learning
-
-    def _folded_weight(self) -> Optional[torch.Tensor]:
-        # language=rst
-        """
-        Return the product of the pipeline's foldable feature values. If the
-        pipeline is empty or contains a non-foldable feature, returns none.
-        """
-        if not self.pipeline:
-            return None
-        if self._w_eff is not None:
-            return self._w_eff
-        w_eff = None
-        for f in self.pipeline:
-            v = f.matmul_fold_value()
-            if v is None:
-                return None  # a non-foldable feature -> use the generic path
-            w_eff = v if w_eff is None else w_eff * v
-        # A product of only boolean masks would be non-float; force the matmul
-        # dtype (float32 by default, or the opt-in low-precision compute_dtype).
-        if self.compute_dtype is not None:
-            w_eff = w_eff.to(self.compute_dtype)
-        elif not torch.is_floating_point(w_eff):
-            w_eff = w_eff.float()
-        if not self.manual_update:
-            self._w_eff = w_eff
-        return w_eff
-
-    def _sparse_matmul(self, s_flat: torch.Tensor, w_eff: torch.Tensor) -> torch.Tensor:
-        # language=rst
-        """
-        Activity-sparse form of ``s_flat @ w_eff``. Read only
-        the rows of ``w_eff`` for source neurons that spiked this step. Numerically
-        identical to the dense matmul. More efficient when few source neurons are active.
-        """
-        active = s_flat.any(dim=0).nonzero(as_tuple=False).squeeze(1)
-        if active.numel() == 0:
-            return torch.zeros(
-                s_flat.size(0), self.target.n, device=w_eff.device, dtype=w_eff.dtype
-            )
-        return s_flat[:, active].to(w_eff.dtype) @ w_eff.index_select(0, active)
-
     def compute(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
         """
-        Compute pre-activations given spikes using connection weights.
+        Direct incoming spikes through the connection's feature pipeline.
 
-        :param s: Incoming spikes.
-        :return: Incoming spikes multiplied by synaptic weights (with or without
-                 decaying spike activation).
-        """
+        Each feature's ``compute`` returns its ``[source.n, target.n]`` value; how
+        it folds is set by the feature's ``op`` (``"mul"`` default, ``"add"``,
+        ``"sub"``). Folding the recurrence (start ``A = 1``, ``B = 0``):
 
-        # Fast path: when the whole pipeline folds to a static elementwise-multiply
-        w_eff = self._folded_weight()
-        if w_eff is not None:
-            s_flat = s.view(s.size(0), self.source.n)
-            if self.sparse_compute:
-                out_signal = self._sparse_matmul(s_flat, w_eff)
-            else:
-                out_signal = s_flat.to(w_eff.dtype) @ w_eff
-            if out_signal.dtype != torch.float32:
-                out_signal = out_signal.float()
-            if self.traces:
-                self.activity = out_signal
-            if out_signal.size() != torch.Size([s.size(0)] + self.target.shape):
-                return out_signal.view(s.size(0), *self.target.shape)
-            return out_signal
-
-        # Change to numeric type (torch doesn't like booleans for matrix ops)
-        # Note: .float() is an expensive operation. Use as minimally as possible!
-        # if s.dtype != torch.float32:
-        #     s = s.float()
-
-        # Prepare broadcast from incoming spikes to all output neurons
-        # |conn_spikes| = [batch_size, source.n * target.n]
-        conn_spikes = s.view(s.size(0), self.source.n, 1).repeat(1, 1, self.target.n)
-        # TODO: ^ This could probably be optimized
-
-        # Run through pipeline
-        for f in self.pipeline:
-            conn_spikes = f.compute(conn_spikes)
+        * ``mul`` factor ``a``: ``A <- a * A`` and ``B <- a * B``
+        * ``add`` term  ``b``: ``B <- B + b``
+        * ``sub`` term  ``b``: ``B <- B - b``
+
+        ``B`` stays ``None`` (unallocated) unless an additive feature is present,
+        so a purely multiplicative pipeline is exactly the single ``s @ A`` matmul.
 
-        # Sum signals for each of the output/terminal neurons
-        # |out_signal| = [batch_size, target.n]
-        if conn_spikes.size() != torch.Size([s.size(0), self.source.n, self.target.n]):
-            if conn_spikes.is_sparse:
-                conn_spikes = conn_spikes.to_dense()
-            conn_spikes = conn_spikes.view(s.size(0), self.source.n, self.target.n)
+        :param s: Incoming spikes, shape ``[batch, *source.shape]``.
+        :return: Post-synaptic input, shape ``[batch, *target.shape]``.
+        """
+        s = s.view(s.size(0), self.source.n)
 
-        if conn_spikes.is_sparse:
-            out_signal = conn_spikes.to_dense().sum(1)
+        # running product of multiplicative factors, [source.n, target.n]
+        a_eff = None
+        # running additive offset, [source.n, target.n]; None while still zero
+        b_eff = None
+        for f in self.pipeline:
+            factor = f.compute(s)
+            op = getattr(f, "op", "mul")
+            if op == "mul":
+                a_eff = factor if a_eff is None else a_eff * factor
+                if b_eff is not None:
+                    b_eff = b_eff * factor
+            else:  # additive contribution: "add" -> +factor, "sub" -> -factor
+                term = factor if op == "add" else -factor
+                b_eff = term if b_eff is None else b_eff + term
+
+        if a_eff is None:
+            # Degenerate pipeline with no multiplicative feature: every source
+            # neuron contributes with unit weight.
+            a_eff = torch.ones(s.size(1), b_eff.size(-1), device=s.device)
+        if not torch.is_floating_point(a_eff):
+            a_eff = a_eff.float()
+
+        if self.sparse_compute:
+            # Read only the rows of A for source neurons that spiked this step
+            # (numerically identical; faster only when few are active).
+            active = s.any(dim=0).nonzero(as_tuple=False).squeeze(1)
+            if active.numel() == 0:
+                out = torch.zeros(
+                    s.size(0), a_eff.size(-1), device=a_eff.device, dtype=a_eff.dtype
+                )
+            else:
+                out = s[:, active].to(a_eff.dtype) @ a_eff.index_select(0, active)
         else:
-            out_signal = conn_spikes.sum(1)
+            out = s.to(a_eff.dtype) @ a_eff
 
-        if self.traces:
-            self.activity = out_signal
+        # Additive terms apply to every synapse regardless of spikes, so sum over
+        # all source rows (the closed form of the [batch, source.n, target.n]
+        # source-sum).
+        if b_eff is not None:
+            out = out + b_eff.sum(dim=0)
 
-        if out_signal.size() != torch.Size([s.size(0)] + self.target.shape):
-            return out_signal.view(s.size(0), *self.target.shape)
-        else:
-            return out_signal
+        if self.traces:
+            self.activity = out
+        if out.size() != torch.Size([s.size(0)] + self.target.shape):
+            return out.view(s.size(0), *self.target.shape)
+        return out
 
     def compute_window(self, s: torch.Tensor) -> torch.Tensor:
         # language=rst
@@ -596,8 +545,6 @@ def update(self, **kwargs) -> None:
             # Pipeline learning
             for f in self.pipeline:
                 f.update(**kwargs)
-            if self._pipeline_has_learning():
-                self._w_eff = None  # weights changed -> rebuild fold next compute
 
     def normalize(self) -> None:
         # language=rst
@@ -607,7 +554,6 @@ def normalize(self) -> None:
         # Normalize pipeline features
         for f in self.pipeline:
             f.normalize()
-        self._w_eff = None  # normalization may change weights -> invalidate fold
 
     def reset_state_variables(self) -> None:
         # language=rst
@@ -618,7 +564,6 @@ def reset_state_variables(self) -> None:
 
         for f in self.pipeline:
             f.reset_state_variables()
-        self._w_eff = None  # rebuild the fold cache at the next sample
 
 
 class Conv1dConnection(AbstractConnection):
diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py
index c39288020..f6126eba2 100644
--- a/bindsnet/network/topology_features.py
+++ b/bindsnet/network/topology_features.py
@@ -18,6 +18,12 @@ class AbstractFeature(ABC):
     Features to operate on signals traversing a connection.
     """
 
+    # How this feature folds in the connection's affine pipeline (see
+    # :meth:`MulticompartmentConnection.compute`): ``"mul"`` (elementwise multiply,
+    # the default), ``"add"`` or ``"sub"`` (elementwise add/subtract of the value
+    # returned by :meth:`compute`).
+    op = "mul"
+
     @abstractmethod
     def __init__(
         self,
@@ -163,28 +169,16 @@ def reset_state_variables(self) -> None:
         pass
 
     @abstractmethod
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
         # language=rst
         """
-        Computes the feature being operated on a set of incoming signals.
+        Return this feature's ``[source.n, target.n]`` value, given pre-synaptic
+        spikes ``s`` of shape ``[batch, source.n]``. How the value folds into the
+        connection is set by :attr:`op` -- a multiplicative factor (``"mul"``, the
+        default), or an additive/subtractive offset (``"add"``/``"sub"``).
         """
         pass
 
-    def matmul_fold_value(self) -> Optional[torch.Tensor]:
-        # language=rst
-        """
-        Fast-path hook for :class:`MulticompartmentConnection`.
-
-        If this feature is a *pure elementwise multiply* by a static
-        ``[source.n, target.n]`` tensor -- i.e. ``compute(x) == x * V`` with no
-        per-call side effects on ``V`` -- return that ``V`` so the connection can
-        fold the whole pipeline into one weight matrix and evaluate as a matmul
-        (``out = s @ W_eff``) instead of materialising ``[B, source.n,
-        target.n]``. Returning ``None`` (the safe default) forces the generic
-        pipeline path.
-        """
-        return None
-
     def prime_feature(self, connection, device, **kwargs) -> None:
         # language=rst
         """
@@ -437,11 +431,11 @@ def sparse_bernoulli(self):
         non_zero = values[mask]
         return torch.sparse_coo_tensor(indices, non_zero, self.value.size())
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        # Factor: a fresh Bernoulli draw each step (resampled, never cached).
         if self.sparse:
-            return conn_spikes * self.sparse_bernoulli()
-        else:
-            return conn_spikes * torch.bernoulli(self.value)
+            return self.sparse_bernoulli()
+        return torch.bernoulli(self.value)
 
     def reset_state_variables(self) -> None:
         pass
@@ -519,13 +513,7 @@ def __init__(
         self.name = name
         self.value = value
 
-    def compute(self, conn_spikes) -> torch.Tensor:
-        return conn_spikes * self.value
-
-    def matmul_fold_value(self) -> Optional[torch.Tensor]:
-        # A boolean mask is a pure elementwise multiply (True->1, False->0).
-        if getattr(self, "sparse", False):
-            return None
+    def compute(self, s) -> torch.Tensor:
         return self.value
 
     def reset_state_variables(self) -> None:
@@ -581,9 +569,9 @@ def __init__(self) -> None:
     def reset_state_variables(self) -> None:
         pass
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
-        return conn_spikes.mean() * torch.ones(
-            self.source_n * self.target_n, device=self.device
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        return s.float().mean() * torch.ones(
+            self.source_n, self.target_n, device=s.device
         )
 
     def prime_feature(self, connection, device, **kwargs) -> None:
@@ -651,7 +639,7 @@ def __init__(
     def reset_state_variables(self) -> None:
         pass
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
         if self.enforce_polarity:
             pos_mask = ~torch.logical_xor(self.value > 0, self.positive_mask)
             neg_mask = ~torch.logical_xor(self.value < 0, ~self.positive_mask)
@@ -659,19 +647,11 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
             self.value[~pos_mask] = 0.0001
             self.value[~neg_mask] = -0.0001
 
-        return_val = self.value * conn_spikes
+        factor = self.value
         if self.norm_frequency == "time step":
+            factor = factor.clone()
             self.normalize(time_step_norm=True)
-
-        return return_val
-
-    def matmul_fold_value(self) -> Optional[torch.Tensor]:
-        # Foldable only as a plain multiply: skip when polarity enforcement or
-        # per-time-step normalisation would mutate ``value`` during compute, or
-        # when the value is stored sparse.
-        if self.enforce_polarity or self.norm_frequency == "time step" or self.sparse:
-            return None
-        return self.value
+        return factor
 
     def prime_feature(self, connection, device, **kwargs) -> None:
         #### Initialize value ####
@@ -734,11 +714,15 @@ def __init__(
             batch_size=batch_size,
         )
 
+    # Bias is additive: folds as ``B <- B + value`` in the connection's pipeline.
+    op = "add"
+
     def reset_state_variables(self) -> None:
         pass
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
-        return conn_spikes + self.value
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        # Additive offset added to every synapse (independent of the spikes).
+        return self.value
 
     def prime_feature(self, connection, device, **kwargs) -> None:
         #### Initialize value ####
@@ -762,7 +746,7 @@ def __init__(
     ) -> None:
         # language=rst
         """
-        Adds scalars to signals
+        Multiply all signals by a scalar
         :param name: Name of the feature
         :param value: Values to scale signals by
         :param value_dtype: Data type for :code:`value` tensor
@@ -781,8 +765,8 @@ def __init__(
     def reset_state_variables(self) -> None:
         pass
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
-        return conn_spikes * self.value
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        return self.value
 
     def prime_feature(self, connection, device, **kwargs) -> None:
         #### Initialize value ####
@@ -835,11 +819,17 @@ def __init__(
 
         self.degrade_function = degrade_function
 
+    # Degradation is subtractive: folded as ``B <- B - degrade_function(value)``.
+    op = "sub"
+
     def reset_state_variables(self) -> None:
         pass
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
-        return conn_spikes - self.degrade_function(self.value)
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        # Subtractive offset (via degrade_function) applied to every synapse.
+        if self.degrade_function is not None:
+            return self.degrade_function(self.value)
+        return self.value
 
 
 class AdaptationBaseSynapsHistory(AbstractFeature):
@@ -911,7 +901,11 @@ def forward(self, x):
             batch_size=batch_size,
         )
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        # This feature needs the per-synapse spikes, so build them from s
+        conn_spikes = s.view(s.size(0), -1, 1).expand(
+            s.size(0), s.size(1), self.value.size(-1)
+        )
 
         # Update the spike buffer
         if self.start_counter == False or conn_spikes.sum() > 0:
@@ -940,7 +934,7 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
             if self.sparse:
                 self.value = self.value.to_sparse()
 
-        return conn_spikes * self.value
+        return self.value
 
     def reset_state_variables(
         self,
@@ -1021,7 +1015,11 @@ def forward(self, x):
             batch_size=batch_size,
         )
 
-    def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
+    def compute(self, s) -> Union[torch.Tensor, float, int]:
+        # This feature needs the per-synapse spikes, so build them from s
+        conn_spikes = s.view(s.size(0), -1, 1).expand(
+            s.size(0), s.size(1), self.value.size(-1)
+        )
 
         # Update the spike buffer
         if self.start_counter == False or conn_spikes.sum() > 0:
@@ -1050,7 +1048,7 @@ def compute(self, conn_spikes) -> Union[torch.Tensor, float, int]:
             if self.sparse:
                 self.value = self.value.to_sparse()
 
-        return conn_spikes * self.value
+        return self.value
 
     def reset_state_variables(
         self,
@@ -1089,15 +1087,16 @@ def __init__(
         self.parent = parent_feature
         self.sub_feature = None  # <-- Defined in non-abstract constructor
 
-    def compute(self, _) -> None:
+    def compute(self, s):
         # language=rst
         """
-        Proxy function to catch a pipeline execution from topology.py's :code:`compute` function. Allows :code:`SubFeature`
-        objects to be executed like real features in the pipeline.
+        Proxy to run a parent feature's side-effect (e.g. normalize/update) from
+        inside the pipeline. Returns ``1`` so it is the identity factor in the fold.
         """
 
         # sub_feature should be defined in the non-abstract constructor
         self.sub_feature()
+        return 1
 
 
 class Normalization(AbstractSubFeature):

From 48c98c04bed38b65ffeaa8104c3aa97c1ef4b92e Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Sun, 26 Jul 2026 20:21:56 +0900
Subject: [PATCH 6/8] Fixed indentation error with learning rule class

---
 bindsnet/learning/MCC_learning.py | 137 +++++++++++++++---------------
 1 file changed, 69 insertions(+), 68 deletions(-)

diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py
index 3fe2014c9..660bd538c 100644
--- a/bindsnet/learning/MCC_learning.py
+++ b/bindsnet/learning/MCC_learning.py
@@ -304,89 +304,90 @@ def _connection_update(self, **kwargs) -> None:
     def reset_state_variables(self):
         return
 
-    class Hebbian(MCC_LearningRule):
+
+class Hebbian(MCC_LearningRule):
+    # language=rst
+    """
+    Simple Hebbian learning rule. Pre- and post-synaptic updates are both positive.
+    """
+
+    def __init__(
+        self,
+        connection: AbstractMulticompartmentConnection,
+        feature_value: Union[torch.Tensor, float, int],
+        nu: Optional[Union[float, Sequence[float]]] = None,
+        reduction: Optional[callable] = None,
+        decay: float = 0.0,
+        **kwargs,
+    ) -> None:
         # language=rst
         """
-        Simple Hebbian learning rule. Pre- and post-synaptic updates are both positive.
-        """
+        Constructor for ``Hebbian`` learning rule.
 
-        def __init__(
-            self,
-            connection: AbstractMulticompartmentConnection,
-            feature_value: Union[torch.Tensor, float, int],
-            nu: Optional[Union[float, Sequence[float]]] = None,
-            reduction: Optional[callable] = None,
-            decay: float = 0.0,
+        :param connection: An ``AbstractConnection`` object whose weights the
+            ``Hebbian`` learning rule will modify.
+        :param nu: Single or pair of learning rates for pre- and post-synaptic events.
+        :param reduction: Method for reducing parameter updates along the batch
+            dimension.
+        :param decay: Coefficient controlling rate of decay of the weights each iteration.
+        """
+        super().__init__(
+            connection=connection,
+            feature_value=feature_value,
+            nu=nu,
+            reduction=reduction,
+            decay=decay,
             **kwargs,
-        ) -> None:
-            # language=rst
-            """
-            Constructor for ``Hebbian`` learning rule.
-
-            :param connection: An ``AbstractConnection`` object whose weights the
-                ``Hebbian`` learning rule will modify.
-            :param nu: Single or pair of learning rates for pre- and post-synaptic events.
-            :param reduction: Method for reducing parameter updates along the batch
-                dimension.
-            :param decay: Coefficient controlling rate of decay of the weights each iteration.
-            """
-            super().__init__(
-                connection=connection,
-                feature_value=feature_value,
-                nu=nu,
-                reduction=reduction,
-                decay=decay,
-                **kwargs,
-            )
+        )
 
-            assert (
-                self.source.traces and self.target.traces
-            ), "Both pre- and post-synaptic nodes must record spike traces."
+        assert (
+            self.source.traces and self.target.traces
+        ), "Both pre- and post-synaptic nodes must record spike traces."
 
-            if isinstance(MulticompartmentConnection):
-                self.update = self._connection_update
-                self.feature_value = feature_value
-            # elif isinstance(connection, Conv2dConnection):
-            #     self.update = self._conv2d_connection_update
-            else:
-                raise NotImplementedError(
-                    "This learning rule is not supported for this Connection type."
-                )
+        if isinstance(MulticompartmentConnection):
+            self.update = self._connection_update
+            self.feature_value = feature_value
+        # elif isinstance(connection, Conv2dConnection):
+        #     self.update = self._conv2d_connection_update
+        else:
+            raise NotImplementedError(
+                "This learning rule is not supported for this Connection type."
+            )
 
-        def _connection_update(self, **kwargs) -> None:
-            # language=rst
-            """
-            Hebbian learning rule for ``Connection`` subclass of ``AbstractConnection``
-            class.
-            """
+    def _connection_update(self, **kwargs) -> None:
+        # language=rst
+        """
+        Hebbian learning rule for ``Connection`` subclass of ``AbstractConnection``
+        class.
+        """
 
-            # Add polarities back to feature after updates
-            if self.enforce_polarity:
-                self.feature_value = torch.abs(self.feature_value)
+        # Add polarities back to feature after updates
+        if self.enforce_polarity:
+            self.feature_value = torch.abs(self.feature_value)
 
-            batch_size = self.source.batch_size
+        batch_size = self.source.batch_size
 
-            source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float()
-            source_x = self.source.x.view(batch_size, -1).unsqueeze(2)
-            target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float()
-            target_x = self.target.x.view(batch_size, -1).unsqueeze(1)
+        source_s = self.source.s.view(batch_size, -1).unsqueeze(2).float()
+        source_x = self.source.x.view(batch_size, -1).unsqueeze(2)
+        target_s = self.target.s.view(batch_size, -1).unsqueeze(1).float()
+        target_x = self.target.x.view(batch_size, -1).unsqueeze(1)
 
-            # Pre-synaptic update.
-            update = self.reduction(torch.bmm(source_s, target_x), dim=0)
-            self.feature_value += self.nu[0] * update
+        # Pre-synaptic update.
+        update = self.reduction(torch.bmm(source_s, target_x), dim=0)
+        self.feature_value += self.nu[0] * update
 
-            # Post-synaptic update.
-            update = self.reduction(torch.bmm(source_x, target_s), dim=0)
-            self.feature_value += self.nu[1] * update
+        # Post-synaptic update.
+        update = self.reduction(torch.bmm(source_x, target_s), dim=0)
+        self.feature_value += self.nu[1] * update
 
-            # Add polarities back to feature after updates
-            if self.enforce_polarity:
-                self.feature_value = self.feature_value * self.polarities
+        # Add polarities back to feature after updates
+        if self.enforce_polarity:
+            self.feature_value = self.feature_value * self.polarities
 
-            super().update()
+        super().update()
 
-        def reset_state_variables(self):
-            return
+    def reset_state_variables(self):
+        return
 
 
 class MSTDP(MCC_LearningRule):

From 6f3c5a68151c745574a292c654a392805b27c6e7 Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Sun, 26 Jul 2026 21:45:25 +0900
Subject: [PATCH 7/8] Test cases and minor bug fixes

---
 bindsnet/learning/MCC_learning.py     |   4 +-
 bindsnet/network/topology_features.py |   2 +
 test/network/test_connections.py      | 408 ++++++++++++++++++++++++++
 3 files changed, 413 insertions(+), 1 deletion(-)

diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py
index 660bd538c..9c2a7b781 100644
--- a/bindsnet/learning/MCC_learning.py
+++ b/bindsnet/learning/MCC_learning.py
@@ -315,6 +315,7 @@ def __init__(
         self,
         connection: AbstractMulticompartmentConnection,
         feature_value: Union[torch.Tensor, float, int],
+        range: Optional[Sequence[float]] = None,
         nu: Optional[Union[float, Sequence[float]]] = None,
         reduction: Optional[callable] = None,
         decay: float = 0.0,
@@ -334,6 +335,7 @@ def __init__(
         super().__init__(
             connection=connection,
             feature_value=feature_value,
+            range=[-1, +1] if range is None else range,
             nu=nu,
             reduction=reduction,
             decay=decay,
@@ -344,7 +346,7 @@ def __init__(
             self.source.traces and self.target.traces
         ), "Both pre- and post-synaptic nodes must record spike traces."
 
-        if isinstance(MulticompartmentConnection):
+        if isinstance(connection, MulticompartmentConnection):
             self.update = self._connection_update
             self.feature_value = feature_value
         # elif isinstance(connection, Conv2dConnection):
diff --git a/bindsnet/network/topology_features.py b/bindsnet/network/topology_features.py
index f6126eba2..3cfc5a672 100644
--- a/bindsnet/network/topology_features.py
+++ b/bindsnet/network/topology_features.py
@@ -84,6 +84,7 @@ def __init__(
         from ..learning.MCC_learning import (
             NoOp,
             PostPre,
+            Hebbian,
             MSTDP,
             MSTDPET,
         )
@@ -91,6 +92,7 @@ def __init__(
         supported_rules = [
             NoOp,
             PostPre,
+            Hebbian,
             MSTDP,
             MSTDPET,
         ]
diff --git a/test/network/test_connections.py b/test/network/test_connections.py
index 24f2333e8..416855359 100644
--- a/test/network/test_connections.py
+++ b/test/network/test_connections.py
@@ -1,4 +1,5 @@
 import torch
+import math
 
 from bindsnet.learning import (
     MSTDP,
@@ -12,6 +13,8 @@
 from bindsnet.network import Network
 from bindsnet.network.nodes import Input, LIFNodes, SRM0Nodes
 from bindsnet.network.topology import *
+import bindsnet.learning.MCC_learning as mcc
+import bindsnet.network.topology_features as tf
 
 
 class TestConnection:
@@ -164,7 +167,412 @@ def test_weights(self, conn_type, shape_a, shape_b, shape_w, *args, **kwargs):
                     )
 
 
+class TestMultiCompartmentConnection:
+
+    def __init__(self):
+        self.device = torch.device("cpu")
+        print(f"Using device '{self.device}' for the MCC test")
+
+    # ----------------------------------------------------------------------- #
+    # Helpers                                                                 #
+    # ----------------------------------------------------------------------- #
+
+    def _make_mcc(self, pipeline, src_n, tgt_n, batch=1, sparse_compute=False):
+        """Build (and prime) a standalone MCC: Input(src_n) -> LIFNodes(tgt_n)."""
+        src = Input(n=src_n, traces=True)
+        tgt = LIFNodes(n=tgt_n, traces=True)
+        # batch_size is None until a Network sets it; the learning rules need it.
+        src.batch_size = batch
+        tgt.batch_size = batch
+        conn = MulticompartmentConnection(
+            source=src,
+            target=tgt,
+            device=self.device,
+            pipeline=pipeline,
+            sparse_compute=sparse_compute,
+        )
+        conn.dt = 1.0  # not set until added to a Network; rules read connection.dt
+        return conn
+
+    def _reference_expansion(self, pipeline, s, tgt_n):
+        """Pre-collapse pipeline features """
+        b, src = s.shape
+        x = s.view(b, src, 1).expand(b, src, tgt_n).clone().float()
+        for f in pipeline:
+            op = getattr(f, "op", "mul")
+            if isinstance(f, tf.Degradation):
+                v = (
+                    f.degrade_function(f.value)
+                    if f.degrade_function is not None
+                    else f.value
+                )
+            else:
+                v = f.value
+            if torch.is_tensor(v):
+                v = v.float()
+            if op == "mul":
+                x = x * v
+            elif op == "add":
+                x = x + v
+            else:  # "sub"
+                x = x - v
+        return x.sum(dim=1)
+
+    def _learning_conn(self, rule, w0, nu, rng=(-1.0, 1.0)):
+        """MCC with a single learnable Weight; returns (connection, feature)."""
+        src_n, tgt_n = w0.shape
+        conn = self._make_mcc(
+            [tf.Weight(name="w", value=w0.clone(), learning_rule=rule, nu=nu, range=rng)],
+            src_n,
+            tgt_n,
+            batch=1,
+        )
+        return conn, conn.pipeline[0]
+
+    def _mstdp_seq(self):
+        return [
+            (torch.tensor([1.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 0.0),
+            (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([1.0, 0.0]), 1.0),
+            (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 1.0),
+            (torch.tensor([0.0, 0.0, 0.0]), torch.tensor([0.0, 0.0]), 1.0),
+        ]
+
+    def _mstdp_reference(self, w0, seq, nu0, dt, tc_plus=20.0, tc_minus=20.0, rng=(-1.0, 1.0)):
+        w = w0.clone().float()
+        src_n, tgt_n = w.shape
+        p_plus, p_minus = torch.zeros(src_n), torch.zeros(tgt_n)
+        elig = torch.zeros(src_n, tgt_n)
+        dp, dm = math.exp(-dt / tc_plus), math.exp(-dt / tc_minus)
+        for src_s, tgt_s, reward in seq:
+            w = w + nu0 * reward * elig
+            p_plus = dp * p_plus + src_s
+            p_minus = dm * p_minus - tgt_s
+            elig = torch.outer(p_plus, tgt_s) + torch.outer(src_s, p_minus)
+            w = torch.clamp(w, rng[0], rng[1])
+        return w
+
+    def _mstdpet_reference(
+        self, w0, seq, nu0, dt, tc_plus=20.0, tc_minus=20.0, tc_e=25.0, rng=(-1.0, 1.0)
+    ):
+        w = w0.clone().float()
+        src_n, tgt_n = w.shape
+        p_plus, p_minus = torch.zeros(src_n), torch.zeros(tgt_n)
+        elig = torch.zeros(src_n, tgt_n)
+        elig_tr = torch.zeros(src_n, tgt_n)
+        for src_s, tgt_s, reward in seq:
+            elig_tr = elig_tr * math.exp(-dt / tc_e) + elig / tc_e
+            w = w + nu0 * dt * reward * elig_tr
+            p_plus = p_plus * math.exp(-dt / tc_plus) + src_s
+            p_minus = p_minus * math.exp(-dt / tc_minus) - tgt_s
+            elig = torch.outer(p_plus, tgt_s) + torch.outer(src_s, p_minus)
+            w = torch.clamp(w, rng[0], rng[1])
+        return w
+
+    # ----------------------------------------------------------------------- #
+    # Individual feature outputs                                              #
+    # ----------------------------------------------------------------------- #
+
+    def test_weight_feature_output(self):
+        s = torch.tensor([[1.0, 0.0, 1.0], [0.0, 1.0, 1.0]])
+        w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        conn = self._make_mcc([tf.Weight(name="w", value=w.clone())], 3, 2, batch=2)
+        assert torch.allclose(conn.compute(s), s @ w, atol=1e-6)
+
+    def test_mask_feature_output(self):
+        s = torch.tensor([[1.0, 1.0, 1.0]])
+        w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        mask = torch.tensor([[True, False], [True, True], [False, True]])
+        conn = self._make_mcc(
+            [tf.Weight(name="w", value=w.clone()), tf.Mask(name="m", value=mask.clone())],
+            3,
+            2,
+        )
+        assert torch.allclose(conn.compute(s), s @ (w * mask.float()), atol=1e-6)
+
+    def test_bias_feature_output(self):
+        # Bias is additive per-synapse; after the source-sum it adds bias.sum(0).
+        s = torch.tensor([[1.0, 0.0, 1.0], [1.0, 1.0, 0.0]])
+        w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        bias = torch.tensor([[0.5, -0.5], [0.1, 0.2], [0.0, 1.0]])
+        conn = self._make_mcc(
+            [tf.Weight(name="w", value=w.clone()), tf.Bias(name="b", value=bias.clone())],
+            3,
+            2,
+            batch=2,
+        )
+        assert torch.allclose(conn.compute(s), s @ w + bias.sum(0), atol=1e-6)
+
+    def test_intensity_feature_output(self):
+        # Intensity's value is a per-synapse [src, tgt] tensor (a constant 2.0 here).
+        s = torch.tensor([[1.0, 0.0, 1.0]])
+        w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        intensity = torch.full((3, 2), 2.0)
+        conn = self._make_mcc(
+            [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Intensity(name="i", value=intensity.clone(), range=(-5.0, 5.0)),
+            ],
+            3,
+            2,
+        )
+        assert torch.allclose(conn.compute(s), s @ (w * intensity), atol=1e-6)
+
+    def test_degradation_feature_output(self):
+        # Degradation subtracts degrade_function(value) per-synapse -> -sum(0).
+        s = torch.tensor([[1.0, 1.0, 0.0]])
+        w = torch.tensor([[0.4, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        deg = torch.tensor([[0.2, 0.4], [0.6, 0.8], [0.1, 0.3]])
+        conn = self._make_mcc(
+            [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Degradation(
+                    name="d", value=deg.clone(), degrade_function=lambda v: v * 0.5
+                ),
+            ],
+            3,
+            2,
+        )
+        assert torch.allclose(conn.compute(s), s @ w - (0.5 * deg).sum(0), atol=1e-6)
+
+    def test_probability_feature_deterministic_bounds(self):
+        # bernoulli(1) == 1 (always passes); bernoulli(0) == 0 (always blocked).
+        s = torch.tensor([[1.0, 0.0, 1.0], [1.0, 1.0, 1.0]])
+        w = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.5, 0.6]])
+        passes = self._make_mcc(
+            [
+                tf.Probability(name="p", value=torch.ones(3, 2)),
+                tf.Weight(name="w", value=w.clone()),
+            ],
+            3,
+            2,
+            batch=2,
+        )
+        assert torch.allclose(passes.compute(s), s @ w, atol=1e-6)
+        blocked = self._make_mcc(
+            [
+                tf.Probability(name="p", value=torch.zeros(3, 2)),
+                tf.Weight(name="w", value=w.clone()),
+            ],
+            3,
+            2,
+            batch=2,
+        )
+        assert torch.allclose(blocked.compute(s), torch.zeros(2, 2), atol=1e-6)
+
+    def test_adaptation_features_output(self):
+        for cls in (tf.AdaptationBaseSynapsHistory, tf.AdaptationBaseOtherSynaps):
+            src_n, tgt_n = 3, 2
+            s = torch.tensor([[1.0, 0.0, 1.0]])
+            feat = cls(
+                name="a",
+                value=torch.zeros(src_n, tgt_n),
+                ann_values=[torch.zeros(1, 1), torch.zeros(1, 1)],
+            )
+            conn = self._make_mcc([feat], src_n, tgt_n, batch=1)
+            out = conn.compute(s)
+            assert torch.allclose(feat.value.float(), torch.ones(src_n, tgt_n))
+            assert torch.allclose(out, s @ torch.ones(src_n, tgt_n), atol=1e-6)
+
+    # ----------------------------------------------------------------------- #
+    # Combined pipelines == pre-collapse expansion                            #
+    # ----------------------------------------------------------------------- #
+
+    def test_pipeline_matches_expansion(self):
+        torch.manual_seed(0)
+        src_n, tgt_n, batch = 4, 3, 2
+        s = (torch.rand(batch, src_n) > 0.4).float()
+        w = torch.randn(src_n, tgt_n)
+        w2 = torch.randn(src_n, tgt_n)
+        mask = torch.rand(src_n, tgt_n) > 0.5
+        bias = torch.randn(src_n, tgt_n) * 0.2
+        deg = torch.rand(src_n, tgt_n)
+
+        pipelines = {
+            "weight": [tf.Weight(name="w", value=w.clone())],
+            "weight+mask": [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Mask(name="m", value=mask.clone()),
+            ],
+            "weight+bias": [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Bias(name="b", value=bias.clone()),
+            ],
+            "weight+bias+degradation": [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Bias(name="b", value=bias.clone()),
+                tf.Degradation(
+                    name="d", value=deg.clone(), degrade_function=lambda v: v * 0.3
+                ),
+            ],
+            "weight+intensity+bias": [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Intensity(
+                    name="i", value=torch.full((src_n, tgt_n), 1.5), range=(-5.0, 5.0)
+                ),
+                tf.Bias(name="b", value=bias.clone()),
+            ],
+            "weight,bias,weight,bias": [
+                tf.Weight(name="w", value=w.clone()),
+                tf.Bias(name="b", value=bias.clone()),
+                tf.Weight(name="w2", value=w2.clone()),
+                tf.Bias(name="b2", value=(bias * 0.5).clone()),
+            ],
+        }
+
+        for name, pipe in pipelines.items():
+            conn = self._make_mcc(pipe, src_n, tgt_n, batch=batch)
+            out = conn.compute(s)
+            ref = self._reference_expansion(pipe, s, tgt_n)
+            assert torch.allclose(
+                out, ref, atol=1e-5
+            ), f"{name}: {(out - ref).abs().max()}"
+
+    # ----------------------------------------------------------------------- #
+    # Sparse activity / sparse_compute == dense                               #
+    # ----------------------------------------------------------------------- #
+
+    def test_sparse_compute_matches_dense(self):
+        torch.manual_seed(1)
+        src_n, tgt_n, batch = 6, 4, 3
+        w = torch.randn(src_n, tgt_n)
+        bias = torch.randn(src_n, tgt_n) * 0.2
+        deg = torch.rand(src_n, tgt_n)
+
+        def pipes():
+            return [
+                [tf.Weight(name="w", value=w.clone())],
+                [
+                    tf.Weight(name="w", value=w.clone()),
+                    tf.Bias(name="b", value=bias.clone()),
+                ],
+                [
+                    tf.Weight(name="w", value=w.clone()),
+                    tf.Degradation(
+                        name="d", value=deg.clone(), degrade_function=lambda v: v * 0.4
+                    ),
+                ],
+            ]
+
+        spike_sets = [
+            (torch.rand(batch, src_n) > 0.5).float(),
+            torch.zeros(batch, src_n),  # empty-spike edge case
+        ]
+        for dense_pipe, sparse_pipe in zip(pipes(), pipes()):
+            for s in spike_sets:
+                dense = self._make_mcc(dense_pipe, src_n, tgt_n, batch=batch).compute(s)
+                sparse = self._make_mcc(
+                    sparse_pipe, src_n, tgt_n, batch=batch, sparse_compute=True
+                ).compute(s)
+                assert torch.allclose(dense, sparse, atol=1e-5)
+
+    # ----------------------------------------------------------------------- #
+    # MCC learning rules                                                      #
+    # ----------------------------------------------------------------------- #
+
+    def test_noop_leaves_weight_unchanged(self):
+        w0 = torch.rand(3, 2)
+        conn, feat = self._learning_conn(mcc.NoOp, w0, nu=(0.1, 0.1))
+        conn.source.s = torch.ones(1, 3)
+        conn.target.s = torch.ones(1, 2)
+        feat.learning_rule.update(reward=1.0)
+        assert torch.allclose(feat.value, w0)
+
+    def test_postpre_predictable(self):
+        # PostPre: dW = -nu0 * outer(src_s, tgt_x) + nu1 * outer(src_x, tgt_s).
+        w0 = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.0, -0.2]])
+        conn, feat = self._learning_conn(mcc.PostPre, w0, nu=(0.1, 0.2))
+        conn.source.s = torch.tensor([[1.0, 0.0, 1.0]])
+        conn.target.s = torch.tensor([[1.0, 1.0]])
+        conn.source.x = torch.tensor([[0.3, 0.7, 0.2]])
+        conn.target.x = torch.tensor([[0.5, 0.4]])
+        feat.learning_rule.update()
+        dw = -0.1 * torch.outer(conn.source.s[0], conn.target.x[0]) + 0.2 * torch.outer(
+            conn.source.x[0], conn.target.s[0]
+        )
+        expected = torch.clamp(w0 + dw, -1.0, 1.0)
+        assert torch.allclose(feat.value, expected, atol=1e-6)
+        assert not torch.allclose(feat.value, w0)  # sanity: it actually changed
+
+    def test_learning_respects_range_clamp(self):
+        # Post-only potentiation of +5 per synapse must clamp to the range max (1.0).
+        w0 = torch.full((2, 2), 0.9)
+        conn, feat = self._learning_conn(mcc.PostPre, w0, nu=(0.0, 5.0), rng=(-1.0, 1.0))
+        conn.source.s = torch.zeros(1, 2)
+        conn.target.s = torch.ones(1, 2)
+        conn.source.x = torch.ones(1, 2)
+        conn.target.x = torch.zeros(1, 2)
+        feat.learning_rule.update()
+        assert torch.allclose(feat.value, torch.ones(2, 2))
+
+    def test_mstdp_predictable(self):
+        w0 = torch.zeros(3, 2)
+        conn, feat = self._learning_conn(mcc.MSTDP, w0, nu=(0.5, 0.0))
+        seq = self._mstdp_seq()
+        for src_v, tgt_v, reward in seq:
+            conn.source.s = src_v.unsqueeze(0)
+            conn.target.s = tgt_v.unsqueeze(0)
+            feat.learning_rule.update(reward=reward)
+        expected = self._mstdp_reference(w0, seq, nu0=0.5, dt=1.0)
+        assert torch.allclose(feat.value, expected, atol=1e-5)
+        assert not torch.allclose(feat.value, w0)
+
+    def test_mstdpet_predictable(self):
+        w0 = torch.zeros(3, 2)
+        conn, feat = self._learning_conn(mcc.MSTDPET, w0, nu=(0.5, 0.5))
+        seq = self._mstdp_seq()
+        for src_v, tgt_v, reward in seq:
+            conn.source.s = src_v.unsqueeze(0)
+            conn.target.s = tgt_v.unsqueeze(0)
+            feat.learning_rule.update(reward=reward)
+        expected = self._mstdpet_reference(w0, seq, nu0=0.5, dt=1.0)
+        assert torch.allclose(feat.value, expected, atol=1e-5)
+        assert not torch.allclose(feat.value, w0)
+
+    def test_hebbian_predictable(self):
+        # Hebbian: dW = nu0 * outer(src_s, tgt_x) + nu1 * outer(src_x, tgt_s), then
+        # clamped. Both pre- and post-synaptic terms are positive (contrast PostPre,
+        # which subtracts the pre term).
+        w0 = torch.tensor([[0.1, 0.2], [0.3, 0.4], [0.0, -0.2]])
+        conn, feat = self._learning_conn(mcc.Hebbian, w0, nu=(0.1, 0.2))
+        conn.source.s = torch.tensor([[1.0, 0.0, 1.0]])
+        conn.target.s = torch.tensor([[1.0, 1.0]])
+        conn.source.x = torch.tensor([[0.3, 0.7, 0.2]])
+        conn.target.x = torch.tensor([[0.5, 0.4]])
+        feat.learning_rule.update()
+        dw = 0.1 * torch.outer(conn.source.s[0], conn.target.x[0]) + 0.2 * torch.outer(
+            conn.source.x[0], conn.target.s[0]
+        )
+        expected = torch.clamp(w0 + dw, -1.0, 1.0)
+        assert torch.allclose(feat.value, expected, atol=1e-6)
+        assert not torch.allclose(feat.value, w0)  # sanity: it actually changed
+
+
 if __name__ == "__main__":
+    # MulticompartmentConnection + MCC learning-rule tests.
+    mcc_tester = TestMultiCompartmentConnection()
+    mcc_tests = [
+        mcc_tester.test_weight_feature_output,
+        mcc_tester.test_mask_feature_output,
+        mcc_tester.test_bias_feature_output,
+        mcc_tester.test_intensity_feature_output,
+        mcc_tester.test_degradation_feature_output,
+        mcc_tester.test_probability_feature_deterministic_bounds,
+        mcc_tester.test_meanfield_not_implemented,
+        mcc_tester.test_adaptation_features_output,
+        mcc_tester.test_pipeline_matches_expansion,
+        mcc_tester.test_sparse_compute_matches_dense,
+        mcc_tester.test_noop_leaves_weight_unchanged,
+        mcc_tester.test_postpre_predictable,
+        mcc_tester.test_learning_respects_range_clamp,
+        mcc_tester.test_mstdp_predictable,
+        mcc_tester.test_mstdpet_predictable,
+        mcc_tester.test_hebbian_predictable,
+    ]
+    for mcc_test in mcc_tests:
+        mcc_test()
+        print(f"  PASSED: {mcc_test.__name__}")
+    print(f"All {len(mcc_tests)} MulticompartmentConnection tests passed.")
+
     tester = TestConnection()
 
     # tester.test_transfer()

From a81c05dbb2b05baa556b7bb4a41d76c10ef4b1e0 Mon Sep 17 00:00:00 2001
From: christopher-earl <cearl@umass.edu>
Date: Sun, 26 Jul 2026 23:04:16 +0900
Subject: [PATCH 8/8] Benchmark tests

---
 .../_bench_common.py                          | 137 +++++++++++++
 .../foldable_pipelines.py                     |   7 +
 .../run_both.py                               |   9 +
 .../sparse_compute.py                         |   8 +
 examples/stress_test/example_network.py       | 185 ++++++++++++++++++
 5 files changed, 346 insertions(+)
 create mode 100644 examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py
 create mode 100644 examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py
 create mode 100644 examples/benchmark/sparse_compute and foldable pipeline/run_both.py
 create mode 100644 examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py
 create mode 100644 examples/stress_test/example_network.py

diff --git a/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py b/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py
new file mode 100644
index 000000000..564de6591
--- /dev/null
+++ b/examples/benchmark/sparse_compute and foldable pipeline/_bench_common.py	
@@ -0,0 +1,137 @@
+"""
+Shared machinery for the ExampleNetwork MCC benchmarks (imported by
+``foldable_pipelines.py``, ``sparse_compute.py`` and ``both.py``).
+
+Each benchmark reports the %-time speedup of one MCC configuration over another by
+timing the *same* ExampleNetwork under different ``compute`` modes:
+
+  * ``expansion``   -- the pre-optimization path, reconstructed here as a
+                       monkeypatch: materialize ``[batch, src, tgt]``, apply every
+                       feature elementwise, then sum over source. (This path was
+                       removed from the code when the fold landed, so we rebuild it
+                       to serve as the baseline.)
+  * ``fold``        -- the current folded path: ``out = s @ A + B.sum(0)``.
+  * ``fold_sparse`` -- the fold plus activity-sparse compute (``sparse_compute``):
+                       read only the weight rows of source neurons that spiked.
+
+Speedup is wall-time reduction: ``(t_baseline - t_new) / t_baseline * 100``
+(positive = faster). Learning is disabled so we time the pure forward path (the
+part these optimizations affect).
+"""
+
+import os
+import statistics
+import sys
+import time
+
+_HERE = os.path.dirname(os.path.abspath(__file__))
+_ROOT = os.path.dirname(os.path.dirname(_HERE))  # repo root (for ``bindsnet``)
+_STRESS = os.path.join(_ROOT, "stress_test")  # for ``example_network``
+for _p in (_ROOT, _STRESS):
+    if _p not in sys.path:
+        sys.path.insert(0, _p)
+
+import torch
+
+from bindsnet.network.topology import MulticompartmentConnection
+from bindsnet.network.topology_features import Degradation, Probability
+from example_network import ExampleNetwork
+
+# ExampleNetwork sizes per device: 20k excitatory neurons on GPU (where the fold
+# shines), a smaller net on CPU so the baseline finishes in reasonable time.
+GPU_CONFIG = dict(in_size=100, exc_size=20_000, inh_size=2_000)
+CPU_CONFIG = dict(in_size=100, exc_size=2_000, inh_size=200)
+GPU_TIME, CPU_TIME = 50, 20
+REPS, WARMUP = 5, 2
+
+
+def _expansion_compute(self, s):
+    """The pre-fold ``[batch, src, tgt]`` expansion path (baseline)."""
+    s = s.view(s.size(0), self.source.n)
+    cs = s.view(s.size(0), self.source.n, 1).repeat(1, 1, self.target.n)
+    for f in self.pipeline:
+        op = getattr(f, "op", "mul")
+        if isinstance(f, Probability):
+            v = torch.bernoulli(f.value)
+        elif isinstance(f, Degradation):
+            v = f.degrade_function(f.value) if f.degrade_function is not None else f.value
+        else:
+            v = f.value
+        if op == "mul":
+            cs = cs * v
+        elif op == "add":
+            cs = cs + v
+        else:
+            cs = cs - v
+    out = cs.sum(1)
+    if getattr(self, "traces", False):
+        self.activity = out
+    if out.size() != torch.Size([s.size(0)] + self.target.shape):
+        return out.view(s.size(0), *self.target.shape)
+    return out
+
+
+def _set_mode(net, mode):
+    for c in net.connections.values():
+        if isinstance(c, MulticompartmentConnection):
+            if mode == "expansion":
+                c.compute = _expansion_compute.__get__(c)  # instance override
+                c.sparse_compute = False
+            else:
+                c.__dict__.pop("compute", None)  # restore the class (fold) method
+                c.sparse_compute = mode == "fold_sparse"
+
+
+def _bench(net, inputs, T, device, modes):
+    cuda = device.startswith("cuda")
+    net.train(False)  # forward-only: time the compute path the optimizations touch
+    for m in modes:  # warmup each mode
+        _set_mode(net, m)
+        for _ in range(WARMUP):
+            net.reset_state_variables()
+            net.run(inputs=inputs, time=T)
+    if cuda:
+        torch.cuda.synchronize()
+    samples = {m: [] for m in modes}
+    for _ in range(REPS):  # interleave modes each rep to counter thermal drift
+        for m in modes:
+            _set_mode(net, m)
+            net.reset_state_variables()
+            if cuda:
+                torch.cuda.synchronize()
+            t0 = time.perf_counter()
+            net.run(inputs=inputs, time=T)
+            if cuda:
+                torch.cuda.synchronize()
+            samples[m].append(time.perf_counter() - t0)
+    return {m: statistics.median(v) * 1e3 for m, v in samples.items()}
+
+
+def run_speedup(baseline_mode, test_mode, technique):
+    """Build the ExampleNetwork on CPU and GPU, time both modes, print the speedup."""
+    print("=" * 72)
+    print(f"{technique}")
+    print(f"  (%-time speedup of '{test_mode}' vs baseline '{baseline_mode}')")
+    print("=" * 72)
+
+    devices = ["cpu"] + (["cuda"] if torch.cuda.is_available() else [])
+    if not torch.cuda.is_available():
+        print("[note] CUDA not available -> reporting CPU only.\n")
+
+    for dev in devices:
+        cfg = GPU_CONFIG if dev == "cuda" else CPU_CONFIG
+        T = GPU_TIME if dev == "cuda" else CPU_TIME
+        net = ExampleNetwork(device=dev, **cfg)
+        inputs = net.make_input(T)
+        res = _bench(net, inputs, T, dev, [baseline_mode, test_mode])
+        base, new = res[baseline_mode], res[test_mode]
+        speedup = (base - new) / base * 100.0
+        print(
+            f"  [{dev.upper():4s}] exc={cfg['exc_size']:>6d}  time={T:>3d}  |  "
+            f"{baseline_mode}={base:8.2f} ms   {test_mode}={new:8.2f} ms   "
+            f"|  speedup = {speedup:+6.1f}% time"
+        )
+        del net
+        if torch.cuda.is_available():
+            torch.cuda.empty_cache()
+    print()
diff --git a/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py b/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py
new file mode 100644
index 000000000..ea43adfbe
--- /dev/null
+++ b/examples/benchmark/sparse_compute and foldable pipeline/foldable_pipelines.py	
@@ -0,0 +1,7 @@
+from _bench_common import run_speedup
+
+### Benchmark: foldable pipelines ###
+# Measures the speedup of the folded MultiCompartmentConnection compute (``out = s @ A + B.sum(0)``) over
+# the pre-fold ``[batch, src, tgt]`` expansion path, on the ExampleNetwork
+if __name__ == "__main__":
+    run_speedup(baseline_mode="expansion", test_mode="fold", technique="Foldable pipelines")
diff --git a/examples/benchmark/sparse_compute and foldable pipeline/run_both.py b/examples/benchmark/sparse_compute and foldable pipeline/run_both.py
new file mode 100644
index 000000000..0752da254
--- /dev/null
+++ b/examples/benchmark/sparse_compute and foldable pipeline/run_both.py	
@@ -0,0 +1,9 @@
+from _bench_common import run_speedup
+
+# Run both sparse_compute.py and foldable_pipelines.py
+if __name__ == "__main__":
+    run_speedup(
+        baseline_mode="expansion",
+        test_mode="fold_sparse",
+        technique="Foldable pipelines + sparse compute (combined)",
+    )
diff --git a/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py b/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py
new file mode 100644
index 000000000..3a4a6f733
--- /dev/null
+++ b/examples/benchmark/sparse_compute and foldable pipeline/sparse_compute.py	
@@ -0,0 +1,8 @@
+from _bench_common import run_speedup
+
+### Benchmark: sparse compute ###
+# Measures the speedup of activity-sparse compute (``sparse_compute=True`` -- read
+# only the weight rows of source neurons that spiked) over the dense folded compute,
+# on the ExampleNetwork (CPU and GPU).
+if __name__ == "__main__":
+    run_speedup(baseline_mode="fold", test_mode="fold_sparse", technique="Sparse compute")
diff --git a/examples/stress_test/example_network.py b/examples/stress_test/example_network.py
new file mode 100644
index 000000000..7562a2736
--- /dev/null
+++ b/examples/stress_test/example_network.py
@@ -0,0 +1,185 @@
+"""
+ExampleNetwork -- a large, sparse, recurrent ``MulticompartmentConnection`` (MCC)
+stress workload.
+
+Topology: ``Input(I) -> EXC_LIF <-> INH_LIF``. Four single-``Weight`` MCC.
+The ``I -> EXC`` connection also carries an ``MSTDP`` learning rule.
+
+Run it directly to stress the simulator:
+
+    python examples/stress_test/example_network.py --device cuda --exc 20000 --time 50
+    python examples/stress_test/example_network.py --device cpu  --exc 2000  --time 20
+"""
+
+import argparse
+import os
+import sys
+import time
+
+# Make ``bindsnet`` importable when this file is run as a standalone script.
+sys.path.insert(
+    0, os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+)
+
+import torch
+
+from bindsnet.learning.MCC_learning import MSTDP
+from bindsnet.network.network import Network
+from bindsnet.network.nodes import Input, LIFNodes
+from bindsnet.network.topology import MulticompartmentConnection
+from bindsnet.network.topology_features import Weight
+
+
+class ExampleNetwork(Network):
+    def __init__(
+        self,
+        device="cpu",
+        in_size=100,
+        exc_size=20_000,
+        inh_size=2_000,
+        batch_size=1,
+        i_to_exc_connectivity=0.15,
+        i_to_inh_connectivity=0.05,
+        inh_to_exc_connectivity=0.05,
+        exc_to_inh_connectivity=0.05,
+    ):
+        super().__init__()
+        self.device = device
+        self.in_size = in_size
+        self.exc_size = exc_size
+        self.inh_size = inh_size
+        self.batch_size = batch_size
+        self.i_to_exc_connectivity = i_to_exc_connectivity
+        self.i_to_inh_connectivity = i_to_inh_connectivity
+        self.inh_to_exc_connectivity = inh_to_exc_connectivity
+        self.exc_to_inh_connectivity = exc_to_inh_connectivity
+        self.build()
+
+    def _sparse_weight(self, rows, cols, connectivity, sign=1.0):
+        w = sign * torch.rand(rows, cols, device=self.device)
+        keep = torch.rand(rows, cols, device=self.device) > (1 - connectivity)
+        return w * keep
+
+    def build(self):
+        device = self.device
+        self.add_layer(layer=Input(self.in_size), name="I")
+        self.add_layer(layer=LIFNodes(self.exc_size), name="EXC_LIF")
+        self.add_layer(layer=LIFNodes(self.inh_size), name="INH_LIF")
+        self.add_connection(
+            connection=MulticompartmentConnection(
+                source=self.layers["I"],
+                target=self.layers["EXC_LIF"],
+                device=device,
+                pipeline=[
+                    Weight(
+                        name="I_to_EXC_weight",
+                        value=self._sparse_weight(
+                            self.in_size, self.exc_size, self.i_to_exc_connectivity
+                        ),
+                        learning_rule=MSTDP,
+                        range=(0, 1),
+                    )
+                ],
+            ),
+            source="I",
+            target="EXC_LIF",
+        )
+        self.add_connection(
+            connection=MulticompartmentConnection(
+                source=self.layers["I"],
+                target=self.layers["INH_LIF"],
+                device=device,
+                pipeline=[
+                    Weight(
+                        name="I_to_INH_weight",
+                        value=self._sparse_weight(
+                            self.in_size, self.inh_size, self.i_to_inh_connectivity
+                        ),
+                    )
+                ],
+            ),
+            source="I",
+            target="INH_LIF",
+        )
+        self.add_connection(
+            connection=MulticompartmentConnection(
+                source=self.layers["INH_LIF"],
+                target=self.layers["EXC_LIF"],
+                device=device,
+                pipeline=[
+                    Weight(
+                        name="INH_to_EXC_weight",
+                        value=self._sparse_weight(
+                            self.inh_size,
+                            self.exc_size,
+                            self.inh_to_exc_connectivity,
+                            sign=-1.0,
+                        ),
+                    )
+                ],
+            ),
+            source="INH_LIF",
+            target="EXC_LIF",
+        )
+        self.add_connection(
+            connection=MulticompartmentConnection(
+                source=self.layers["EXC_LIF"],
+                target=self.layers["INH_LIF"],
+                device=device,
+                pipeline=[
+                    Weight(
+                        name="EXC_to_INH_weight",
+                        value=self._sparse_weight(
+                            self.exc_size, self.inh_size, self.exc_to_inh_connectivity
+                        ),
+                    )
+                ],
+            ),
+            source="EXC_LIF",
+            target="INH_LIF",
+        )
+        self.to(device)
+
+    def make_input(self, runtime):
+        # Poisson-ish random spike train into the input layer.
+        return {
+            "I": torch.rand(
+                runtime, self.batch_size, self.in_size, device=self.device
+            )
+            > 0.90
+        }
+
+
+if __name__ == "__main__":
+    p = argparse.ArgumentParser(description="Stress-test the ExampleNetwork.")
+    p.add_argument("--device", default="cuda" if torch.cuda.is_available() else "cpu")
+    p.add_argument("--in-size", type=int, default=100)
+    p.add_argument("--exc", type=int, default=20_000)
+    p.add_argument("--inh", type=int, default=2_000)
+    p.add_argument("--time", type=int, default=50)
+    args = p.parse_args()
+
+    device = args.device
+    if device.startswith("cuda") and not torch.cuda.is_available():
+        print("[note] CUDA unavailable; falling back to CPU.")
+        device = "cpu"
+
+    net = ExampleNetwork(
+        device=device, in_size=args.in_size, exc_size=args.exc, inh_size=args.inh
+    )
+    net.train(False)  # forward-only stress (no learning)
+    inputs = net.make_input(args.time)
+
+    net.run(inputs=inputs, time=args.time)  # warmup
+    if device.startswith("cuda"):
+        torch.cuda.synchronize()
+    net.reset_state_variables()
+    t0 = time.perf_counter()
+    net.run(inputs=inputs, time=args.time)
+    if device.startswith("cuda"):
+        torch.cuda.synchronize()
+    ms = (time.perf_counter() - t0) * 1e3
+    print(
+        f"ExampleNetwork [{device}] in={args.in_size} exc={args.exc} inh={args.inh} "
+        f"time={args.time}: {ms:.1f} ms total ({ms / args.time:.3f} ms/step)"
+    )