cleaned up graded/patched comps with inner batched_reset formulation

Author: Alexander Ororbia

ngclearn/components/neurons/graded/bernoulliErrorCell.py

@@ -105,13 +105,15 @@ def advance_state(self, dt): ## compute Bernoulli error cell output
         self.L.set(jnp.squeeze(L))
         self.mask.set(mask)
 
+    @compilable
+    def reset(self):  ## reset core components/statistics
+        self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
 
-    # @transition(output_compartments=["dp", "dtarget", "target", "p", "modulator", "L", "mask"])
     @compilable
-    def reset(self): ## reset core components/statistics
-        _shape = (self.batch_size, self.shape[0])
+    def batched_reset(self, batch_size):
+        _shape = (batch_size, self.shape[0])
         if len(self.shape) > 1:
-            _shape = (self.batch_size, self.shape[0], self.shape[1], self.shape[2])
+            _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape) ## "rest"/reset values
         dp = restVals
         dtarget = restVals
@@ -130,7 +132,6 @@ def reset(self): ## reset core components/statistics
         self.L.set(L)
         self.mask.set(mask)
 
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -47,6 +47,7 @@ def __init__(self, name, n_units, batch_size=1, sigma=1., shape=None, **kwargs):
             sigma_shape = jnp.array(sigma).shape
         self.sigma_shape = sigma_shape
         self.shape = shape
+        self.batch_size = batch_size
         self.n_units = n_units
 
         ## Convolution shape setup
@@ -104,10 +105,12 @@ def advance_state(self, dt): ## compute Gaussian error cell output
         self.L.set(jnp.squeeze(L))
         self.mask.set(mask)
 
-    # @transition(output_compartments=["dmu", "dtarget", "dSigma", "target", "mu", "modulator", "L", "mask"])
-    # @staticmethod
     @compilable
-    def reset(self, batch_size): ## reset core components/statistics
+    def reset(self): ## reset core components/statistics
+        self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size):
         _shape = (batch_size, self.shape[0])
         if len(self.shape) > 1:
             _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -100,16 +100,20 @@ def advance_state(self, dt): ## compute Laplacian error cell output
         self.mask.set(mask)
 
     @compilable
-    def reset(self): ## reset core components/statistics
-        restVals = jnp.zeros((self.batch_size, self.n_units))
+    def reset(self):  ## reset core components/statistics
+        self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size):
+        restVals = jnp.zeros((batch_size, self.n_units))
         dshift = restVals
         dtarget = restVals
         dScale = jnp.zeros(self.scale_shape)
         target = restVals
         shift = restVals
         modulator = shift + 1.
         L = 0.
-        mask = jnp.ones((self.batch_size, self.n_units))
+        mask = jnp.ones((batch_size, self.n_units))
 
         self.dshift.set(dshift)
         self.dtarget.set(dtarget)

ngclearn/components/neurons/graded/leakyNoiseCell.py

@@ -1,5 +1,6 @@
 from jax import numpy as jnp, random, jit
-from ngcsimlib.logger import info
+#from ngcsimlib.logger import info
+from ngcsimlib import deprecate_args
 from ngclearn import compilable #from ngcsimlib.parser import compilable
 from ngclearn import Compartment #from ngcsimlib.compartment import Compartment
 from ngclearn.components.jaxComponent import JaxComponent
@@ -65,6 +66,7 @@ class LeakyNoiseCell(JaxComponent): ## Real-valued, leaky noise cell
         leak_scale: degree to which membrane leak should be scaled (Default: 1)
     """
 
+    @deprecate_args(sigma_rec="sigma_pre")
     def __init__(
             self, name, n_units, tau_x, act_fx="relu", integration_type="euler", batch_size=1, sigma_pre=0.1,
             sigma_post=0.1, leak_scale=1., shape=None, **kwargs
@@ -127,13 +129,19 @@ def advance_state(self, t, dt):
         self.r_prime.set(r_prime)
 
     @compilable
-    def reset(self):
-        _shape = (self.batch_size, self.shape[0])
+    def reset(self):  ## reset core components/statistics
+        self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size):
+        _shape = (batch_size, self.shape[0])
         if len(self.shape) > 1:
-            _shape = (self.batch_size, self.shape[0], self.shape[1], self.shape[2])
+            _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape)
-        self.j_input.set(restVals)
-        self.j_recurrent.set(restVals)
+        if not self.j_input.targeted:
+            self.j_input.set(restVals)
+        if not self.j_recurrent.targeted:
+            self.j_recurrent.set(restVals)
         self.x.set(restVals)
         self.r.set(restVals)
         self.r_prime.set(restVals)

ngclearn/components/neurons/graded/rateCell.py

@@ -252,7 +252,11 @@ def advance_state(self, dt):
         self.zF.set(zF)
 
     @compilable
-    def reset(self, batch_size):
+    def reset(self):  ## reset core components/statistics
+        self.batched_reset(batch_size=self.batch_size)  ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size):
         _shape = (batch_size, self.shape[0])
         if len(self.shape) > 1:
             _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])

ngclearn/components/synapses/patched/hebbianPatchedSynapse.py

@@ -281,20 +281,24 @@ def evolve(self):
         self.dBiases.set(dBiases)
 
     @compilable
-    def reset(self, batch_size):
+    def reset(self):  ## closed, no-batch argument reset
+        ## write reset command to call inner batched_reset command
+        self.batched_reset(batch_size=self.batch_size) ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size): ## open, batch argument reset
         preVals = jnp.zeros((batch_size, self.shape[0]))
         postVals = jnp.zeros((batch_size, self.shape[1]))
         # BUG: the self.inputs here does not have the targeted field
         # NOTE: Quick workaround is to check if targeted is in the input or not
-        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals) # inputs
-        self.outputs.set(postVals) # outputs
-        self.post_in.set(postVals) # post_in
-        self.pre_out.set(preVals) # pre_out
-        self.pre.set(preVals) # pre
-        self.post.set(postVals) # post
-        self.dWeights.set(jnp.zeros(self.shape)) # dW
-        self.dBiases.set(jnp.zeros(self.shape[1])) # db
-
+        hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals)  # inputs
+        self.outputs.set(postVals)  # outputs
+        self.post_in.set(postVals)  # post_in
+        self.pre_out.set(preVals)  # pre_out
+        self.pre.set(preVals)  # pre
+        self.post.set(postVals)  # post
+        self.dWeights.set(jnp.zeros(self.shape))  # dW
+        self.dBiases.set(jnp.zeros(self.shape[1]))  # db
 
     @classmethod
     def help(cls): ## component help function

ngclearn/components/synapses/patched/patchedSynapse.py

@@ -167,10 +167,15 @@ def advance_state(self):
         self.pre_out.set(pre_out)
 
     @compilable
-    def reset(self, batch_size):
+    def reset(self): ## closed, no-batch argument reset
+        ## write reset command to call inner batched_reset command
+        self.batched_reset(batch_size=self.batch_size) ## arg = batch_size data-member
+
+    @compilable
+    def batched_reset(self, batch_size): ## open, batch argument reset
         preVals = jnp.zeros((batch_size, self.shape[0]))
         postVals = jnp.zeros((batch_size, self.shape[1]))
-        
+
         # BUG: the self.inputs here does not have the targeted field
         # NOTE: Quick workaround is to check if targeted is in the input or not
         hasattr(self.inputs, "targeted") and not self.inputs.targeted and self.inputs.set(preVals)

tests/components/neurons/graded/test_leakyNoiseCell.py

@@ -16,7 +16,14 @@ def test_LeakyNoiseCell1():
   dt = 1.  # ms
   with Context(name) as ctx:
     a = LeakyNoiseCell(
-        name="a", n_units=1, tau_x=50., act_fx="identity", integration_type="euler", batch_size=1, sigma_rec=0.,
+        name="a",
+        n_units=1,
+        tau_x=50.,
+        act_fx="identity",
+        integration_type="euler",
+        batch_size=1,
+        sigma_pre=0.,
+        sigma_post=0.,
         leak_scale=0.
     )
     advance_process = (MethodProcess("advance_proc") >> a.advance_state)

tests/components/synapses/patched/test_patchedSynapse.py

@@ -36,6 +36,7 @@ def test_patchedSynapse():
 
     advance_process = (MethodProcess("advance_proc") >> a.advance_state)
     reset_process = (MethodProcess("reset_proc") >> a.reset)
+    #batch_reset_process = (MethodProcess("batched_reset_proc") >> a.batched_reset)
 
     def clamp_inputs(x):
       a.inputs.set(x)
@@ -46,6 +47,7 @@ def clamp_inputs(x):
   expected_outputs = (jnp.matmul(inputs_seq, weights) * resist_scale) + biases
   outputs_outs = []
   reset_process.run()
+  #batch_reset_process.run(batch_size=batch_size)
   clamp_inputs(inputs_seq)
   advance_process.run(t=0., dt=dt)
   outputs_outs.append(a.outputs.get())