Fixed an execution bug

Author: willgebhardt

ngclearn/components/jaxComponent.py

@@ -1,6 +1,6 @@
 import time
 
-from typing import Union
+from typing import Union, Dict, Any
 import jax
 from jax import numpy as jnp
 from jax import random
@@ -25,7 +25,6 @@ def __init__(self, name: str, key: Union[jax.Array, None] = None):
         self.key = Compartment(
             random.PRNGKey(time.time_ns()) if key is None else key)
 
-
     def save(self, directory: str):
         """
         The default save method for JaxComponents, it stores the values of all

ngclearn/components/synapses/hebbian/traceSTDPSynapse.py

@@ -83,7 +83,7 @@ def __init__(
         if weight_mask is None:
             self.weight_mask = jnp.ones((1, 1))
         else:
-            self.weight_mask = self.weight_mask
+            self.weight_mask = weight_mask
 
         self.weights.set(self.weights.get() * self.weight_mask)
 
@@ -95,27 +95,27 @@ def __init__(
         self.preTrace = Compartment(preVals)
         self.postTrace = Compartment(postVals)
         self.dWeights = Compartment(self.weights.get() * 0)
-        self.eta = jnp.ones((1, 1)) * eta ## global learning rate
+        self.eta = eta ## global learning rate
 
     def _compute_update(self):
         if self.mu > 0.:
             post_shift = jnp.power(self.w_bound - self.weights.get(), self.mu)
             pre_shift = jnp.power(self.weights.get(), self.mu)
-            dWpost = (post_shift * jnp.matmul((self.preSpike.get() - self.preTrace_target).T, self.postSpike.get())) * self.Aplus
+            dWpost = (post_shift * jnp.matmul((self.preTrace.get() - self.preTrace_target).T, self.postSpike.get())) * self.Aplus
 
             if self.Aminus > 0.:
                 dWpre = -(pre_shift * jnp.matmul(self.preSpike.get().T, self.postTrace.get())) * self.Aminus
             else:
                 dWpre = 0.
 
         else:
-            dWpost = jnp.matmul((self.preSpike.get() - self.preTrace_target).T, self.postSpike.get() * self.Aplus)
+            dWpost = jnp.matmul((self.preTrace.get() - self.preTrace_target).T, self.postSpike.get() * self.Aplus)
             if self.Aminus > 0.:
                 dWpre = -jnp.matmul(self.preSpike.get().T, self.postTrace.get() * self.Aminus)
             else:
                 dWpre = 0.
 
-        dW = (dWpost - dWpre)
+        dW = (dWpost + dWpre)
         return dW
 
     @compilable