refactor bernoulli, laplacian, and rewarderror cells

Author: rxng8

ngclearn/components/neurons/graded/bernoulliErrorCell.py

@@ -3,6 +3,7 @@
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
 from ngclearn.utils.model_utils import sigmoid, d_sigmoid
+from ngcsimlib.compilers.process import transition
 
 class BernoulliErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -58,8 +59,9 @@ def __init__(self, name, n_units, batch_size=1, input_logits=False, shape=None,
         self.modulator = Compartment(restVals + 1.0) # to be set/consumed
         self.mask = Compartment(restVals + 1.0)
 
+    @transition(output_compartments=["dp", "dtarget", "L", "mask"])
     @staticmethod
-    def _advance_state(dt, p, target, modulator, mask, input_logits): ## compute Bernoulli error cell output
+    def advance_state(dt, p, target, modulator, mask, input_logits): ## compute Bernoulli error cell output
         # Moves Bernoulli error cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional
         eps = 0.0001
@@ -89,15 +91,9 @@ def _advance_state(dt, p, target, modulator, mask, input_logits): ## compute Ber
         mask = mask * 0. + 1. ## "eat" the mask as it should only apply at time t
         return dp, dtarget, jnp.squeeze(L), mask
 
-    @resolver(_advance_state)
-    def advance_state(self, dp, dtarget, L, mask):
-        self.dp.set(dp)
-        self.dtarget.set(dtarget)
-        self.L.set(L)
-        self.mask.set(mask)
-
+    @transition(output_compartments=["dp", "dtarget", "target", "p", "modulator", "L", "mask"])
     @staticmethod
-    def _reset(batch_size, shape): ## reset core components/statistics
+    def reset(batch_size, shape): ## reset core components/statistics
         _shape = (batch_size, shape[0])
         if len(shape) > 1:
             _shape = (batch_size, shape[0], shape[1], shape[2])
@@ -111,16 +107,6 @@ def _reset(batch_size, shape): ## reset core components/statistics
         mask = jnp.ones(_shape) ## reset mask
         return dp, dtarget, target, p, modulator, L, mask
 
-    @resolver(_reset)
-    def reset(self, dp, dtarget, target, p, modulator, L, mask):
-        self.dp.set(dp)
-        self.dtarget.set(dtarget)
-        self.target.set(target)
-        self.p.set(p)
-        self.modulator.set(modulator)
-        self.L.set(L)
-        self.mask.set(mask)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/graded/laplacianErrorCell.py

@@ -2,6 +2,7 @@
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
 from ngclearn.utils import tensorstats
+from ngcsimlib.compilers.process import transition
 
 class LaplacianErrorCell(JaxComponent): ## Rate-coded/real-valued error unit/cell
     """
@@ -66,8 +67,9 @@ def __init__(self, name, n_units, batch_size=1, scale=1., shape=None, **kwargs):
         self.modulator = Compartment(restVals + 1.0) ## to be set/consumed
         self.mask = Compartment(restVals + 1.0)
 
+    @transition(output_compartments=["dshift", "dtarget", "dScale", "L", "mask"])
     @staticmethod
-    def _advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplacian error cell output
+    def advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplacian error cell output
         # Moves Laplacian cell dynamics one step forward. Specifically, this routine emulates the error unit
         # behavior of the local cost functional:
         # FIXME: Currently, below does: L(targ, shift) = -||targ - shift||_1/scale
@@ -85,16 +87,9 @@ def _advance_state(dt, shift, target, Scale, modulator, mask): ## compute Laplac
         mask = mask * 0. + 1.  ## "eat" the mask as it should only apply at time t
         return dshift, dtarget, dScale, jnp.squeeze(L), mask
 
-    @resolver(_advance_state)
-    def advance_state(self, dshift, dtarget, dScale, L, mask):
-        self.dshift.set(dshift)
-        self.dtarget.set(dtarget)
-        self.dScale.set(dScale)
-        self.L.set(L)
-        self.mask.set(mask)
-
+    @transition(output_compartments=["dshift", "dtarget", "dScale", "target", "shift", "modulator", "L", "mask"])
     @staticmethod
-    def _reset(batch_size, n_units, scale_shape):
+    def reset(batch_size, n_units, scale_shape):
         restVals = jnp.zeros((batch_size, n_units))
         dshift = restVals
         dtarget = restVals
@@ -106,17 +101,6 @@ def _reset(batch_size, n_units, scale_shape):
         mask = jnp.ones((batch_size, n_units))
         return dshift, dtarget, dScale, target, shift, modulator, L, mask
 
-    @resolver(_reset)
-    def reset(self, dshift, dtarget, dScale, target, shift, modulator, L, mask):
-        self.dshift.set(dshift)
-        self.dtarget.set(dtarget)
-        self.dScale.set(dScale)
-        self.target.set(target)
-        self.shift.set(shift)
-        self.modulator.set(modulator)
-        self.L.set(L)
-        self.mask.set(mask)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

ngclearn/components/neurons/graded/rewardErrorCell.py

@@ -1,6 +1,7 @@
 from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
 from jax import numpy as jnp, jit
+from ngcsimlib.compilers.process import transition
 from ngclearn.utils import tensorstats
 
 class RewardErrorCell(JaxComponent): ## Reward prediction error cell
@@ -50,8 +51,9 @@ def __init__(self, name, n_units, alpha, ema_window_len=10,
         self.accum_reward = Compartment(restVals)  ## accumulated reward signal(s)
         self.n_ep_steps = Compartment(jnp.zeros((self.batch_size, 1))) ## number of episode steps taken
 
+    @transition(output_compartments=["mu", "rpe", "n_ep_steps", "accum_reward"])
     @staticmethod
-    def _advance_state(dt, use_online_predictor, alpha, mu, rpe, reward,
+    def advance_state(dt, use_online_predictor, alpha, mu, rpe, reward,
                        n_ep_steps, accum_reward):
         ## compute/update RPE and predictor values
         accum_reward = accum_reward + reward
@@ -61,41 +63,26 @@ def _advance_state(dt, use_online_predictor, alpha, mu, rpe, reward,
         n_ep_steps = n_ep_steps + 1
         return mu, rpe, n_ep_steps, accum_reward
 
-    @resolver(_advance_state)
-    def advance_state(self, mu, rpe, n_ep_steps, accum_reward):
-        self.mu.set(mu)
-        self.rpe.set(rpe)
-        self.n_ep_steps.set(n_ep_steps)
-        self.accum_reward.set(accum_reward)
-
+    @transition(output_compartments=["mu"])
     @staticmethod
-    def _evolve(dt, use_online_predictor, ema_window_len, n_ep_steps, mu,
+    def evolve(dt, use_online_predictor, ema_window_len, n_ep_steps, mu,
                 accum_reward):
         if use_online_predictor:
             ## total episodic reward signal
             r = accum_reward/n_ep_steps
             mu = (1. - 1./ema_window_len) * mu + (1./ema_window_len) * r
         return mu
 
-    @resolver(_evolve)
-    def evolve(self, mu):
-        self.mu.set(mu)
-
+    @transition(output_compartments=["mu", "rpe", "accum_reward", "n_ep_steps"])
     @staticmethod
-    def _reset(batch_size, n_units):
+    def reset(batch_size, n_units):
         restVals = jnp.zeros((batch_size, n_units))
         mu = restVals
         rpe = restVals
         accum_reward = restVals
         n_ep_steps = jnp.zeros((batch_size, 1))
         return mu, rpe, accum_reward, n_ep_steps
 
-    @resolver(_reset)
-    def reset(self, mu, rpe, accum_reward, n_ep_steps):
-        self.mu.set(mu)
-        self.rpe.set(rpe)
-        self.accum_reward.set(accum_reward)
-        self.n_ep_steps.set(n_ep_steps)
 
     @classmethod
     def help(cls): ## component help function