diff --git a/bindsnet/learning/MCC_learning.py b/bindsnet/learning/MCC_learning.py index 77bdd78b..ef20d44d 100644 --- a/bindsnet/learning/MCC_learning.py +++ b/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 diff --git a/test/network/test_learning.py b/test/network/test_learning.py index 16903781..69eee849 100644 --- a/test/network/test_learning.py +++ b/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(