minor edits to exp-kernel/wtas-cell

Author: Alexander Ororbia

ngclearn/components/neurons/spiking/WTASCell.py

@@ -50,7 +50,7 @@ class WTASCell(JaxComponent): ## winner-take-all spiking cell
         thr_jitter: scale of uniform jitter to add to initialization of thresholds
     """
 
-    # Define Functions
+    #@deprecate_args(thr_base="thrBase")
     def __init__(
             self, name, n_units, tau_m, resist_m=1., thr_base=0.4, thr_gain=0.002, refract_time=0., thr_jitter=0.05,
             **kwargs

ngclearn/components/other/expKernel.py

@@ -74,13 +74,22 @@ def advance_state(self, t):
 
         s = inputs
         ## update spike time window and corresponding window volume
-        tf, epsp = _apply_kernel(tf, s, t, self.tau_w, self.win_len, krn_start=0,
-                                 krn_end=self.win_len-1) #0:win_len-1)
+        tf, epsp = _apply_kernel(
+            tf, s, t, self.tau_w, self.win_len, krn_start=0, krn_end=self.win_len-1
+        ) #0:win_len-1)
 
         # Update compartments
         self.epsp.set(epsp)
         self.tf.set(tf)
 
+    @compilable
+    def reset(self):
+        restVals = jnp.zeros((self.batch_size, self.n_units)) ## inputs, epsp
+        restTensor = jnp.zeros([self.win_len, self.batch_size, self.n_units], jnp.float32)  ## tf
+        self.inputs.set(restVals)
+        self.epsp.set(restVals)
+        self.tf.set(restTensor)
+
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/synapses/denseSynapse.py

@@ -1,7 +1,6 @@
 from jax import random, numpy as jnp, jit
 from ngclearn.components.jaxComponent import JaxComponent
-from ngclearn.utils import tensorstats
-from ngclearn.utils.weight_distribution import initialize_params
+from ngclearn.utils.distribution_generator import DistributionGenerator
 from ngcsimlib.logger import info
 
 from ngcsimlib.compartment import Compartment
@@ -58,10 +57,13 @@ def __init__(
 
         if self.weight_init is None:
             info(self.name, "is using default weight initializer!")
-            self.weight_init = {"dist": "uniform", "amin": 0.025, "amax": 0.8}
-        weights = initialize_params(subkeys[0], self.weight_init, shape)
+            # self.weight_init = {"dist": "uniform", "amin": 0.025, "amax": 0.8}
+            # weights = initialize_params(subkeys[0], self.weight_init, shape)
+            self.weight_init = DistributionGenerator.uniform(0.025, 0.8)
+        #weights = initialize_params(subkeys[0], self.weight_init, shape)
+        weights = self.weight_init(shape, subkeys[0])
 
-        if 0. < p_conn < 1.: ## only non-zero and <1 probs allowed
+        if 0. < p_conn < 1.: ## Modifier/constraint: only non-zero and <1 probs allowed
             p_mask = random.bernoulli(subkeys[1], p=p_conn, shape=shape)
             weights = weights * p_mask ## sparsify matrix
 
@@ -76,9 +78,10 @@ def __init__(
         if self.bias_init is None:
             info(self.name, "is using default bias value of zero (no bias "
                             "kernel provided)!")
-        self.biases = Compartment(initialize_params(subkeys[2], bias_init,
-                                                    (1, shape[1]))
-                                  if bias_init else 0.0)
+        self.biases = Compartment(self.bias_init((1, shape[1]), subkeys[2]) if bias_init else 0.0)
+        # self.biases = Compartment(initialize_params(subkeys[2], bias_init,
+        #                                             (1, shape[1]))
+        #                           if bias_init else 0.0)
 
     @compilable
     def advance_state(self):