fix: restrict kwta (#45)

* fix: restrict kwta

* feat: weight learning tag

conex/behaviors/neurons/specs.py

@@ -91,22 +91,22 @@ def initialize(self, neurons):
 
     def forward(self, neurons):
         will_spike = neurons.v >= neurons.threshold
-        will_spike_v = will_spike * (neurons.v - neurons.threshold)
+        v_values = neurons.v
 
         dim = 0
         if self.dimension is not None:
-            will_spike_v = will_spike_v.view(self.shape)
+            v_values = v_values.view(self.shape)
             will_spike = will_spike.view(self.shape)
             dim = self.dimension
 
         if (will_spike.sum(axis=dim) <= self.k).all():
             return
 
-        k_values, k_winners_indices = torch.topk(
-            will_spike_v, self.k, dim=dim, sorted=False
+        _, k_winners_indices = torch.topk(
+            v_values, self.k, dim=dim, sorted=False
         )
-        min_values = k_values.min(dim=0).values
-        winners = will_spike_v >= min_values.expand(will_spike_v.size())
-        ignored = will_spike * (~winners)
+
+        ignored = will_spike
+        ignored.scatter_(dim, k_winners_indices, False)
 
         neurons.v[ignored.view((-1,))] = neurons.v_reset

conex/behaviors/synapses/learning.py

@@ -28,6 +28,10 @@ def no_bound(w, w_min, w_max):
 
 
 class BaseLearning(Behavior):
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.add_tag("weight_learning")
+
     def get_spike_and_trace(self, synapse):
         src_spike = synapse.src.axon.get_spike(synapse.src, synapse.src_delay)
         dst_spike = synapse.dst.axon.get_spike(synapse.dst, synapse.dst_delay)