fixed some minor bugs in rate-coded cells/hebb-syn

Author: Alexander Ororbia

docs/index.rst

@@ -6,8 +6,8 @@ Welcome to ngc-learn's documentation!
 =====================================
 
 **ngc-learn** is a Python library for building, simulating, and analyzing
-biomimetic computational models, arbitrary predictive processing/coding models,
-and spiking neural networks. This toolkit is built on top of 
+biomimetic and NeuroAI computational models, arbitrary predictive processing/coding models,
+spiking neural networks, and general dynamical systems. This toolkit is built on top of 
 `JAX <https://github.com/google/jax>`_ and is distributed under the 3-Clause BSD license.
 
 .. toctree::

ngclearn/components/neurons/graded/bernoulliErrorCell.py

@@ -110,10 +110,10 @@ def advance_state(self, dt): ## compute Bernoulli error cell output
 
     # @transition(output_compartments=["dp", "dtarget", "target", "p", "modulator", "L", "mask"])
     @compilable
-    def reset(self, batch_size): ## reset core components/statistics
-        _shape = (batch_size, self.shape[0])
+    def reset(self): ## reset core components/statistics
+        _shape = (self.batch_size, self.shape[0])
         if len(self.shape) > 1:
-            _shape = (batch_size, self.shape[0], self.shape[1], self.shape[2])
+            _shape = (self.batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape) ## "rest"/reset values
         dp = restVals
         dtarget = restVals
@@ -161,20 +161,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).get())
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar:

ngclearn/components/neurons/graded/gaussianErrorCell.py

@@ -111,14 +111,14 @@ def advance_state(self, dt): ## compute Gaussian error cell output
     # @transition(output_compartments=["dmu", "dtarget", "dSigma", "target", "mu", "modulator", "L", "mask"])
     # @staticmethod
     @compilable
-    def reset(self, batch_size, shape, sigma_shape): ## reset core components/statistics
-        _shape = (batch_size, shape[0])
-        if len(shape) > 1:
-            _shape = (batch_size, shape[0], shape[1], shape[2])
+    def reset(self): ## reset core components/statistics
+        _shape = (self.batch_size, self.shape[0])
+        if len(self.shape) > 1:
+            _shape = (self.batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape)
         dmu = restVals
         dtarget = restVals
-        dSigma = jnp.zeros(sigma_shape)
+        dSigma = jnp.zeros(self.sigma_shape)
         target = restVals
         mu = restVals
         modulator = mu + 1.
@@ -164,20 +164,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).get())
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar:

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -103,16 +103,16 @@ def advance_state(self, dt): ## compute Laplacian error cell output
         self.L.set(jnp.squeeze(L))
         self.mask.set(mask)
 
-    def reset(self, batch_size, n_units, scale_shape):
-        restVals = jnp.zeros((batch_size, n_units))
+    def reset(self): ## reset core components/statistics
+        restVals = jnp.zeros((self.batch_size, self.n_units))
         dshift = restVals
         dtarget = restVals
-        dScale = jnp.zeros(scale_shape)
+        dScale = jnp.zeros(self.scale_shape)
         target = restVals
         shift = restVals
         modulator = shift + 1.
         L = 0.
-        mask = jnp.ones((batch_size, n_units))
+        mask = jnp.ones((self.batch_size, self.n_units))
 
         self.dshift.set(dshift)
         self.dtarget.set(dtarget)
@@ -152,20 +152,6 @@ def help(cls): ## component help function
                 "hyperparameters": hyperparams}
         return info
 
-    def __repr__(self):
-        comps = [varname for varname in dir(self) if isinstance(getattr(self, varname), Compartment)]
-        maxlen = max(len(c) for c in comps) + 5
-        lines = f"[{self.__class__.__name__}] PATH: {self.name}\n"
-        for c in comps:
-            stats = tensorstats(getattr(self, c).get())
-            if stats is not None:
-                line = [f"{k}: {v}" for k, v in stats.items()]
-                line = ", ".join(line)
-            else:
-                line = "None"
-            lines += f"  {f'({c})'.ljust(maxlen)}{line}\n"
-        return lines
-
 if __name__ == '__main__':
     from ngcsimlib.context import Context
     with Context("Bar") as bar:

ngclearn/components/neurons/graded/rateCell.py

@@ -252,40 +252,39 @@ def advance_state(self, dt):
         self.zF.set(zF)
 
     @compilable
-    def reset(self, batch_size, shape): #n_units
-        _shape = (batch_size, shape[0])
-        if len(shape) > 1:
-            _shape = (batch_size, shape[0], shape[1], shape[2])
+    def reset(self): #, batch_size, shape): #n_units
+        _shape = (self.batch_size, self.shape[0])
+        if len(self.shape) > 1:
+            _shape = (self.batch_size, self.shape[0], self.shape[1], self.shape[2])
         restVals = jnp.zeros(_shape)
         self.j.set(restVals)
         self.j_td.set(restVals)
         self.z.set(restVals)
         self.zF.set(restVals)
 
-
-    def save(self, directory, **kwargs):
-        ## do a protected save of constants, depending on whether they are floats or arrays
-        tau_m = (self.tau_m if isinstance(self.tau_m, float)
-                 else jnp.ones([[self.tau_m]]))
-        priorLeakRate = (self.priorLeakRate if isinstance(self.priorLeakRate, float)
-                         else jnp.ones([[self.priorLeakRate]]))
-        resist_scale = (self.resist_scale if isinstance(self.resist_scale, float)
-                        else jnp.ones([[self.resist_scale]]))
-
-        file_name = directory + "/" + self.name + ".npz"
-        jnp.savez(file_name,
-                  tau_m=tau_m, priorLeakRate=priorLeakRate,
-                  resist_scale=resist_scale) #, key=self.key.value)
-
-    def load(self, directory, seeded=False, **kwargs):
-        file_name = directory + "/" + self.name + ".npz"
-        data = jnp.load(file_name)
-        ## constants loaded in
-        self.tau_m = data['tau_m']
-        self.priorLeakRate = data['priorLeakRate']
-        self.resist_scale = data['resist_scale']
-        #if seeded:
-        #    self.key.set(data['key'])
+    # def save(self, directory, **kwargs):
+    #     ## do a protected save of constants, depending on whether they are floats or arrays
+    #     tau_m = (self.tau_m if isinstance(self.tau_m, float)
+    #              else jnp.ones([[self.tau_m]]))
+    #     priorLeakRate = (self.priorLeakRate if isinstance(self.priorLeakRate, float)
+    #                      else jnp.ones([[self.priorLeakRate]]))
+    #     resist_scale = (self.resist_scale if isinstance(self.resist_scale, float)
+    #                     else jnp.ones([[self.resist_scale]]))
+    #
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     jnp.savez(file_name,
+    #               tau_m=tau_m, priorLeakRate=priorLeakRate,
+    #               resist_scale=resist_scale) #, key=self.key.value)
+    #
+    # def load(self, directory, seeded=False, **kwargs):
+    #     file_name = directory + "/" + self.name + ".npz"
+    #     data = jnp.load(file_name)
+    #     ## constants loaded in
+    #     self.tau_m = data['tau_m']
+    #     self.priorLeakRate = data['priorLeakRate']
+    #     self.resist_scale = data['resist_scale']
+    #     #if seeded:
+    #     #    self.key.set(data['key'])
 
     @classmethod
     def help(cls): ## component help function

ngclearn/components/neurons/graded/rewardErrorCell.py

@@ -92,12 +92,12 @@ def evolve(self, dt):
         self.mu.set(mu)
 
     @compilable
-    def reset(self, batch_size, n_units):
-        restVals = jnp.zeros((batch_size, n_units))
+    def reset(self): ## reset core components/statistics
+        restVals = jnp.zeros((self.batch_size, self.n_units))
         mu = restVals
         rpe = restVals
         accum_reward = restVals
-        n_ep_steps = jnp.zeros((batch_size, 1))
+        n_ep_steps = jnp.zeros((self.batch_size, 1))
         self.mu.set(mu)
         self.rpe.set(rpe)
         self.accum_reward.set(accum_reward)

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -254,15 +254,16 @@ def evolve(self):
         self.dBiases.set(dBiases)
 
     @compilable
-    def reset(self, batch_size, shape):
-        preVals = jnp.zeros((batch_size, shape[0]))
-        postVals = jnp.zeros((batch_size, shape[1]))
-        not self.inputs.targeted and self.inputs.set(preVals) # inputs
+    def reset(self): #, batch_size, shape):
+        preVals = jnp.zeros((self.batch_size, self.shape[0]))
+        postVals = jnp.zeros((self.batch_size, self.shape[1]))
+        #not self.inputs.targeted and self.inputs.set(preVals) # inputs
+        self.inputs.set(preVals)
         self.outputs.set(postVals) # outputs
         self.pre.set(preVals) # pre
         self.post.set(postVals) # post
-        self.dWeights.set(jnp.zeros(shape)) # dW
-        self.dBiases.set(jnp.zeros(shape[1])) # db
+        self.dWeights.set(jnp.zeros(self.shape)) # dW
+        self.dBiases.set(jnp.zeros(self.shape[1])) # db
 
     @classmethod
     def help(cls): ## component help function