Added a JaxProcess

Added Jax Process to allow for scanning over the process.

Author: willgebhardt

ngclearn/__init__.py

@@ -36,7 +36,9 @@
 from ngcsimlib.compartment import Compartment
 from ngcsimlib.resolver import resolver
 from ngcsimlib import utils as sim_utils
-from ngcsimlib.compilers.process import Process, transition
+
+from ngclearn.utils.jaxProcess import JaxProcess
+from ngcsimlib.compilers.process import transition, Process
 
 
 from ngcsimlib import configure, preload_modules

ngclearn/components/base_monitor.py

@@ -284,7 +284,7 @@ def load(self, directory, **kwargs):
                 setattr(self, compartment_path, new_comp)
                 self.compartments.append(new_comp.path)
 
-            self._update_resolver()
+            # self._update_resolver()
 
     def make_plot(self, compartment, ax=None, ylabel=None, xlabel=None, title=None, n=None, plot_func=None):
         vals = self.view(compartment)

ngclearn/utils/jaxProcess.py

@@ -0,0 +1,20 @@
+from ngcsimlib.compilers.process import Process
+from jax.lax import scan as _scan
+
+
+class JaxProcess(Process):
+    def scan(self, xs, arg_order=None, compartments_to_monitor=None, save_state=True):
+        if compartments_to_monitor is None:
+            compartments_to_monitor = []
+        if arg_order is None:
+            arg_order = list(self.get_required_args())
+
+        def _pure(current_state, x):
+            v = self.pure(current_state, **{key: value for key, value in zip(arg_order, x)})
+            return v, [v[c.path] for c in compartments_to_monitor]
+
+        vals, stacked = _scan(_pure, init=self.get_required_state(include_special_compartments=True), xs=xs)
+        if save_state:
+            self.updated_modified_state(vals)
+        return stacked
+