update patched synapses and their test cases

Author: rxng8

ngclearn/components/synapses/patched/hebbianPatchedSynapse.py

@@ -5,6 +5,7 @@
 from ngclearn import resolver, Component, Compartment
 from ngclearn.components.synapses.patched import PatchedSynapse
 from ngclearn.utils import tensorstats
+from ngcsimlib.compilers.process import transition
 
 @partial(jit, static_argnums=[3, 4, 5, 6, 7, 8, 9])
 def _calc_update(pre, post, W, w_mask, w_bound, is_nonnegative=True, signVal=1.,
@@ -240,8 +241,9 @@ def _compute_update(w_mask, w_bound, is_nonnegative, sign_value, prior_type, pri
 
         return dW  * jnp.where(0 != jnp.abs(weights), 1, 0) , db
 
+    @transition(output_compartments=["opt_params", "weights", "biases", "dWeights", "dBiases"])
     @staticmethod
-    def _evolve(w_mask, opt, w_bound, is_nonnegative, sign_value, prior_type, prior_lmbda, pre_wght,
+    def evolve(w_mask, 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 = HebbianPatchedSynapse._compute_update(
@@ -258,16 +260,9 @@ def _evolve(w_mask, opt, w_bound, is_nonnegative, sign_value, prior_type, prior_
         weights = _enforce_constraints(weights, w_mask, 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 (
@@ -280,19 +275,6 @@ def _reset(batch_size, shape):
         )
 
 
-
-    @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 = {

ngclearn/components/synapses/patched/patchedSynapse.py

@@ -3,6 +3,7 @@
 from ngclearn import resolver, Component, Compartment
 from ngclearn.components.jaxComponent import JaxComponent
 from ngclearn.utils import tensorstats
+from ngcsimlib.compilers.process import transition
 from ngclearn.utils.weight_distribution import initialize_params
 from ngcsimlib.logger import info
 import math
@@ -135,28 +136,22 @@ def __init__(self, name, shape, n_sub_models, stride_shape=(0,0), w_mask=None, w
         self.biases = Compartment(initialize_params(subkeys[2], bias_init,
                                                     (1, self.shape[1]))
                                   if bias_init else 0.0)
+
+    @transition(output_compartments=["outputs"])
     @staticmethod
-    def _advance_state(Rscale, inputs, weights, biases):
+    def advance_state(Rscale, inputs, weights, biases):
         outputs = (jnp.matmul(inputs, weights) * Rscale) + biases
         return outputs
 
-    @resolver(_advance_state)
-    def advance_state(self, outputs):
-        self.outputs.set(outputs)
-
+    @transition(output_compartments=["inputs", "outputs"])
     @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]))
         inputs = preVals
         outputs = postVals
         return inputs, outputs
 
-    @resolver(_reset)
-    def reset(self, inputs, outputs):
-        self.inputs.set(inputs)
-        self.outputs.set(outputs)
-
     def save(self, directory, **kwargs):
         file_name = directory + "/" + self.name + ".npz"
         if self.bias_init != None:

tests/components/synapses/patched/test_hebbianPatchedSynapse.py

@@ -0,0 +1,88 @@
+# %%
+
+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 HebbianPatchedSynapse
+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_hebbianPatchedSynapse():
+  np.random.seed(42)
+  name = "hebbian_patched_synapse_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+
+  # model hyper
+  shape = (10, 5)
+  n_sub_models = 2
+  stride_shape = (1, 1)
+  batch_size = 1
+  resist_scale = 1.0
+
+  with Context(name) as ctx:
+    a = HebbianPatchedSynapse(
+      name="a", 
+      shape=shape, 
+      n_sub_models=n_sub_models, 
+      stride_shape=stride_shape,
+      resist_scale=resist_scale,
+      batch_size=batch_size
+    )
+
+    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)
+
+  a.weights.set(jnp.ones((12, 12)) * 0.5) 
+
+  in_pre = jnp.ones((10, 12)) * 1.0
+  in_post = jnp.ones((10, 12)) * 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 == (12, 12), ""
+  assert a.weights.value[0, 0] == 0.5, ""
+
+
+
+# test_hebbianPatchedSynapse()

tests/components/synapses/patched/test_patchedSynapse.py

@@ -0,0 +1,69 @@
+# %%
+
+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 PatchedSynapse
+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_patchedSynapse():
+  np.random.seed(42)
+  name = "patched_synapse_ctx"
+  dkey = random.PRNGKey(42)
+  dkey, *subkeys = random.split(dkey, 100)
+  dt = 1.  # ms
+  # model hyper
+  shape = (10, 5)
+  n_sub_models = 2
+  stride_shape = (1, 1)
+  batch_size = 1
+  resist_scale = 1.0
+  with Context(name) as ctx:
+    a = PatchedSynapse(
+      name="a",
+      shape=shape,
+      n_sub_models=n_sub_models,
+      stride_shape=stride_shape,
+      resist_scale=resist_scale,
+      batch_size=batch_size,
+      weight_init={"dist": "gaussian", "std": 0.1},
+      bias_init={"dist": "constant", "value": 0.0}
+    )
+
+    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")
+
+    # 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")
+
+    @Context.dynamicCommand
+    def clamp_inputs(x):
+      a.inputs.set(x)
+
+  inputs_seq = jnp.asarray(np.random.randn(1, 12))
+  weights = a.weights.value
+  biases = a.biases.value
+  expected_outputs = (jnp.matmul(inputs_seq, weights) * resist_scale) + biases
+  outputs_outs = []
+  ctx.reset()
+  ctx.clamp_inputs(inputs_seq)
+  ctx.run(t=0., dt=dt)
+  outputs_outs.append(a.outputs.value)
+  outputs_outs = jnp.concatenate(outputs_outs, axis=1)
+  # Verify outputs match expected values
+  np.testing.assert_allclose(outputs_outs, expected_outputs, atol=1e-5)
+