Fix MSTDP reset crashing when called before the first run

bindsnet/learning/MCC_learning.py

@@ -471,10 +471,13 @@ def __init__(
         self.tc_plus = torch.tensor(kwargs.get("tc_plus", 20.0))
         self.tc_minus = torch.tensor(kwargs.get("tc_minus", 20.0))
 
-        # State the update path fills in lazily: the previous step's spikes,
-        # kept by the fast path for its rank-1 update, and the dense path's
-        # eligibility. None means "not built yet", which is also the state
-        # ``reset_state_variables`` restores.
+        # State the update path fills in lazily, because it needs the batch
+        # size and device that only the first update knows: P+/P-, the previous
+        # step's spikes kept by the fast path for its rank-1 update, and the
+        # dense path's eligibility. None means "not built yet", which is also
+        # the state ``reset_state_variables`` restores.
+        self.p_plus = None
+        self.p_minus = None
         self._prev_source_s = None
         self._prev_target_s = None
         self.eligibility = None
@@ -506,14 +509,14 @@ def _connection_update(self, **kwargs) -> None:
         batch_size = self.source.batch_size
 
         # Initialize eligibility, P^+, and P^-.
-        if not hasattr(self, "p_plus"):
+        if self.p_plus is None:
             self.p_plus = torch.zeros(
                 # batch_size, *self.source.shape, device=self.source.s.device
                 batch_size,
                 self.source.n,
                 device=self.source.s.device,
             )
-        if not hasattr(self, "p_minus"):
+        if self.p_minus is None:
             self.p_minus = torch.zeros(
                 # batch_size, *self.target.shape, device=self.target.s.device
                 batch_size,
@@ -636,10 +639,9 @@ def reset_state_variables(self) -> None:
         starts from the same state as a freshly-built rule.
         """
 
-        if self.eligibility is not None:
-            self.eligibility.zero_()
-        self.p_plus.zero_()
-        self.p_minus.zero_()
+        for state in (self.eligibility, self.p_plus, self.p_minus):
+            if state is not None:
+                state.zero_()
         if self.average_update > 0:
             self.average_buffer.zero_()
             self.average_buffer_index = 0

test/network/test_learning.py

@@ -376,6 +376,15 @@ def test_mstdp_reset_clears_fast_path_spike_lag(self):
         assert rule._prev_source_s is None
         assert rule._prev_target_s is None
 
+    @pytest.mark.parametrize("rule", [mcc.MSTDP, mcc.MSTDPET, mcc.PostPre])
+    def test_reset_before_first_run_does_not_raise(self, rule):
+        # Some of this state is built lazily on the first update, because only
+        # then are the batch size and device known. Resetting a network before
+        # running it must still work.
+        network, rule_obj = self._build(rule)
+        network.reset_state_variables()
+        assert rule_obj is not None
+
     def test_postpre_reset_clears_average_buffers(self):
         # PostPre's reset was a bare ``return``; both buffers survived.
         network, rule = self._build(