minor revisons/updates to hebb/dense syn, metric utils

Author: Alexander Ororbia

ngclearn/components/neurons/graded/rateCell.py

@@ -226,7 +226,7 @@ def advance_state(self, dt):
             ## self.pressure <-- "top-down" expectation / contextual pressure
             ## self.current <-- "bottom-up" data-dependent signal
             dfx_val = self.dfx(z)
-            j = _modulate(j, dfx_val)
+            j = _modulate(j, dfx_val) ## TODO: make this optional (for NGC circuit dynamics)
             j = j * self.resist_scale
             tmp_z = _run_cell(
                 dt, j, j_td, z, self.tau_m, leak_gamma=self.priorLeakRate, integType=self.intgFlag,

ngclearn/components/synapses/denseSynapse.py

@@ -36,14 +36,20 @@ class DenseSynapse(JaxComponent): ## base dense synaptic cable
         p_conn: probability of a connection existing (default: 1.); setting
             this to < 1 and > 0. will result in a sparser synaptic structure
             (lower values yield sparse structure)
+
+        mask: if non-None, a (multiplicative) mask is applied to this synaptic weight matrix
     """
 
     def __init__(
-            self, name, shape, weight_init=None, bias_init=None, resist_scale=1., p_conn=1., batch_size=1, **kwargs
+            self, name, shape, weight_init=None, bias_init=None, resist_scale=1., p_conn=1., mask=None, batch_size=1,
+            **kwargs
     ):
         super().__init__(name, **kwargs)
 
         self.batch_size = batch_size
+        self.mask = 1.
+        if mask is not None:
+            self.mask = mask
 
         ## Synapse meta-parameters
         self.shape = shape
@@ -79,7 +85,9 @@ def __init__(
 
     @compilable
     def advance_state(self):
-        self.outputs.set((jnp.matmul(self.inputs.get(), self.weights.get()) * self.resist_scale) + self.biases.get())
+        weights = self.weights.get()
+        weights = weights * self.mask
+        self.outputs.set((jnp.matmul(self.inputs.get(), weights) * self.resist_scale) + self.biases.get())
 
     @compilable
     def reset(self):

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -86,7 +86,7 @@ def _enforce_constraints(W, w_bound, is_nonnegative=True):
     """
     _W = W
     if w_bound > 0.:
-        if is_nonnegative == True:
+        if is_nonnegative:
             _W = jnp.clip(_W, 0., w_bound)
         else:
             _W = jnp.clip(_W, -w_bound, w_bound)
@@ -173,7 +173,10 @@ def __init__(
             prior=("constant", 0.), w_decay=0., sign_value=1., optim_type="sgd", pre_wght=1., post_wght=1., 
             p_conn=1., resist_scale=1., batch_size=1, **kwargs
     ):
-        super().__init__(name, shape, weight_init, bias_init, resist_scale, p_conn, batch_size=batch_size, **kwargs)
+        super().__init__(
+            name, shape=shape, weight_init=weight_init, bias_init=bias_init, resist_scale=resist_scale, p_conn=p_conn,
+            batch_size=batch_size, **kwargs
+        )
 
         if w_decay > 0.:
             prior = ('l2', w_decay)
@@ -243,19 +246,20 @@ def calc_update(self):
         post = self.post.get()
         weights = self.weights.get()
         biases = self.biases.get()
-        opt_params = self.opt_params.get()
+        #opt_params = self.opt_params.get()
 
         ## calculate synaptic update values
         dWeights, dBiases = HebbianSynapse._compute_update(
-            self.w_bound, self.is_nonnegative, self.sign_value, self.prior_type, self.prior_lmbda, self.pre_wght, self.post_wght,
-            pre, post, weights
+            self.w_bound, self.is_nonnegative, self.sign_value, self.prior_type, self.prior_lmbda, self.pre_wght,
+            self.post_wght, pre, post, weights
         )
 
         self.dWeights.set(dWeights)
         self.dBiases.set(dBiases)
+        #self.opt_params.set(opt_params)
 
     @compilable
-    def evolve(self):
+    def evolve(self, dt):
         # Get the variables
         pre = self.pre.get()
         post = self.post.get()
@@ -268,6 +272,7 @@ def evolve(self):
             self.w_bound, self.is_nonnegative, self.sign_value, self.prior_type, self.prior_lmbda, self.pre_wght, self.post_wght,
             pre, post, weights
         )
+
         ## conduct a step of optimization - get newly evolved synaptic weight value matrix
         if self.bias_init != None:
             opt_params, [weights, biases] = self.opt(opt_params, [weights, biases], [dWeights, dBiases])

ngclearn/utils/metric_utils.py

@@ -308,7 +308,7 @@ def measure_CatNLL(p, x, offset=1e-7, preserve_batch=False):
         nll = jnp.mean(nll)
     return nll #tf.reduce_mean(nll)
 
-@jit
+@partial(jit, static_argnums=[2])
 def measure_RMSE(mu, x, preserve_batch=False):
     """
     Measures root mean squared error (RMSE). Note: If batch is preserved, this returns a column vector where each
@@ -328,7 +328,7 @@ def measure_RMSE(mu, x, preserve_batch=False):
     mse = measure_MSE(mu, x, preserve_batch=preserve_batch)
     return jnp.sqrt(mse) ## sqrt(MSE) is the root-mean-squared-error
 
-@jit
+@partial(jit, static_argnums=[2])
 def measure_MSE(mu, x, preserve_batch=False):
     """
     Measures mean squared error (MSE), or the negative Gaussian log likelihood with variance of 1.0. Note: If batch
@@ -352,7 +352,7 @@ def measure_MSE(mu, x, preserve_batch=False):
         mse = jnp.mean(mse) # this is proper mse
     return mse
 
-@jit
+@partial(jit, static_argnums=[2])
 def measure_MAE(shift, x, preserve_batch=False):
     """
     Measures mean absolute error (MAE), or the negative Laplacian log likelihood with scale of 1.0. Note: If batch
@@ -376,7 +376,7 @@ def measure_MAE(shift, x, preserve_batch=False):
         mae = jnp.mean(mae) # this is proper mae
     return mae
 
-@jit
+@partial(jit, static_argnums=[3])
 def measure_BCE(p, x, offset=1e-7, preserve_batch=False): #1e-10
     """
     Calculates the negative Bernoulli log likelihood or binary cross entropy (BCE). Note: If batch is preserved,