edit to bern-cell

Author: Alexander Ororbia

ngclearn/components/input_encoders/bernoulliCell.py

@@ -6,8 +6,8 @@
 from ngcsimlib.deprecators import deprecate_args
 from ngcsimlib.logger import info, warn
 
-from ngcsimlib.compilers.process import Process, transition
-from ngcsimlib.component import Component
+from ngcsimlib.compilers.process import transition
+#from ngcsimlib.component import Component
 from ngcsimlib.compartment import Compartment
 
 class BernoulliCell(JaxComponent):
@@ -49,7 +49,7 @@ def __init__(self, name, n_units, batch_size=1, **kwargs):
     @staticmethod
     def advance_state(t, key, inputs, tols):
         ## NOTE: should `inputs` be checked if bounded to [0,1]?
-        key, *subkeys = random.split(key, 2)
+        key, *subkeys = random.split(key, 3)
         outputs = random.bernoulli(subkeys[0], p=inputs).astype(jnp.float32)
         # Updates time-of-last-spike (tols) variable:
         # output = s = binary spike vector

tests/components/input_encoders/test_bernoulliCell.py

@@ -25,20 +25,23 @@ def test_bernoulliCell1():
     with Context("Circuit") as ctx:
         a = BernoulliCell(name="a", n_units=1, key=subkeys[0])
 
-        myProcess = (Process()
-                     >> a.advance_state)
+        advance_process = (Process()
+                           >> a.advance_state)
+        ctx.wrap_and_add_command(advance_process.pure, name="run")
 
-        ctx.wrap_and_add_command(myProcess.pure, name="run")
+        reset_process = (Process()
+                         >> a.reset)
+        ctx.wrap_and_add_command(reset_process.pure, name="reset")
 
-        ## create and compile core simulation commands
-        reset_cmd, reset_args = ctx.compile_by_key(
-            a, compile_key="reset"
-        )
-        ctx.add_command(wrap_command(jit(ctx.reset)), name="reset")
-        advance_cmd, advance_args = ctx.compile_by_key(
-            a,compile_key="advance_state"
-        )
-        ctx.add_command(wrap_command(jit(ctx.advance_state)), name="advance")
+        # ## create and compile core simulation commands
+        # reset_cmd, reset_args = ctx.compile_by_key(
+        #     a, compile_key="reset"
+        # )
+        # ctx.add_command(wrap_command(jit(ctx.reset)), name="reset")
+        # advance_cmd, advance_args = ctx.compile_by_key(
+        #     a,compile_key="advance_state"
+        # )
+        # ctx.add_command(wrap_command(jit(ctx.advance_state)), name="advance")
 
 
         ## set up non-compiled utility commands
@@ -50,11 +53,11 @@ def clamp(x):
     x_seq = jnp.asarray([[1., 1., 0., 0., 1.]], dtype=jnp.float32)
 
     outs = []
-    ctx.reset()
+    #ctx.reset()
     for ts in range(x_seq.shape[1]):
         x_t = jnp.array([[x_seq[0,ts]]]) ## get data at time t
         ctx.clamp(x_t)
-        ctx.advance(t=ts*1., dt=1.)
+        ctx.run(t=ts*1.)#, dt=1.)
         outs.append(a.outputs.value)
     outs = jnp.concatenate(outs, axis=1)