push hebbian synapse

Author: rxng8

ngclearn/components/synapses/hebbian/hebbianSynapse.py

@@ -2,6 +2,7 @@
 from functools import partial
 from ngclearn.utils.optim import get_opt_init_fn, get_opt_step_fn
 from ngclearn import resolver, Component, Compartment
+from ngcsimlib.compilers.process import transition
 from ngclearn.components.synapses import DenseSynapse
 from ngclearn.utils import tensorstats
 from ngcsimlib.deprecators import deprecate_args
@@ -216,8 +217,9 @@ def _compute_update(w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda
             post_wght=post_wght)
         return dW, db
 
+    @transition(output_compartments=["opt_params", "weights", "biases", "dWeights", "dBiases"])
     @staticmethod
-    def _evolve(opt, w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda, pre_wght,
+    def evolve(opt, w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda, pre_wght,
                 post_wght, bias_init, pre, post, weights, biases, opt_params):
         ## calculate synaptic update values
         dWeights, dBiases = HebbianSynapse._compute_update(
@@ -234,16 +236,9 @@ def _evolve(opt, w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda, p
         weights = _enforce_constraints(weights, w_bound, is_nonnegative=is_nonnegative)
         return opt_params, weights, biases, dWeights, dBiases
 
-    @resolver(_evolve)
-    def evolve(self, opt_params, weights, biases, dWeights, dBiases):
-        self.opt_params.set(opt_params)
-        self.weights.set(weights)
-        self.biases.set(biases)
-        self.dWeights.set(dWeights)
-        self.dBiases.set(dBiases)
-
+    @transition(output_compartments=["inputs", "outputs", "pre", "post", "dWeights", "dBiases"])
     @staticmethod
-    def _reset(batch_size, shape):
+    def reset(batch_size, shape):
         preVals = jnp.zeros((batch_size, shape[0]))
         postVals = jnp.zeros((batch_size, shape[1]))
         return (
@@ -255,15 +250,6 @@ def _reset(batch_size, shape):
             jnp.zeros(shape[1]), # db
         )
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs, pre, post, dWeights, dBiases):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-        self.pre.set(pre)
-        self.post.set(post)
-        self.dWeights.set(dWeights)
-        self.dBiases.set(dBiases)
-
     @classmethod
     def help(cls): ## component help function
         properties = {

tests/components/synapses/hebbian/test_hebbianSynapse.py

@@ -0,0 +1,85 @@
+# %%
+
+from jax import numpy as jnp, random, jit
+from ngcsimlib.context import Context
+import numpy as np
+np.random.seed(42)
+from ngclearn.components import HebbianSynapse
+from ngcsimlib.compilers import compile_command, wrap_command
+from numpy.testing import assert_array_equal
+
+from ngcsimlib.compilers.process import Process, transition
+from ngcsimlib.component import Component
+from ngcsimlib.compartment import Compartment
+from ngcsimlib.context import Context
+from ngcsimlib.utils.compartment import Get_Compartment_Batch
+
+
+def test_hebbianSynapse():
+  np.random.seed(42)
+  name = "hebbian_synapse_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+
+  # model hyper
+  shape = (10, 5)
+  batch_size = 1
+  resist_scale = 1.0
+
+  with Context(name) as ctx:
+    a = HebbianSynapse(
+      name="a", 
+      shape=shape, 
+      resist_scale=resist_scale,
+      batch_size=batch_size,
+      prior = ("gaussian", 0.01)
+    )
+
+    advance_process = (Process() >> a.advance_state)
+    ctx.wrap_and_add_command(jit(advance_process.pure), name="run")
+    reset_process = (Process() >> a.reset)
+    ctx.wrap_and_add_command(jit(reset_process.pure), name="reset")
+    evolve_process = (Process() >> a.evolve)
+    ctx.wrap_and_add_command(jit(evolve_process.pure), name="evolve")
+
+    # Compile and add commands
+    # reset_cmd, reset_args = ctx.compile_by_key(a, compile_key="reset")
+    # ctx.add_command(wrap_command(jit(reset_cmd)), name="reset")
+    # advance_cmd, advance_args = ctx.compile_by_key(a, compile_key="advance_state")
+    # ctx.add_command(wrap_command(jit(advance_cmd)), name="run")
+    # evolve_cmd, evolve_args = ctx.compile_by_key(a, compile_key="evolve")
+    # ctx.add_command(wrap_command(jit(evolve_cmd)), name="evolve")
+
+    @Context.dynamicCommand
+    def clamp_inputs(x):
+      a.inputs.set(x)
+
+    @Context.dynamicCommand
+    def clamp_pre(x):
+      a.pre.set(x)
+
+    @Context.dynamicCommand
+    def clamp_post(x):
+      a.post.set(x)
+
+  # Test input sequence
+  # Initial weights
+  a.weights.set(jnp.ones((10, 5)) * 0.5) 
+
+  in_pre = jnp.ones((1, 10)) * 1.0
+  in_post = jnp.ones((1, 5)) * 0.75
+
+  ctx.reset()
+  clamp_pre(in_pre)
+  clamp_post(in_post)
+  ctx.run(t=1. * dt, dt=dt)
+  ctx.evolve(t=1. * dt, dt=dt)
+
+  print(a.weights.value)
+
+  # Basic assertions to check learning dynamics
+  assert a.weights.value.shape == (10, 5), ""
+  assert a.weights.value[0, 0] == 0.5, ""
+
+# test_hebbianSynapse()