use for_loop() instead of make_loop()

Author: chaoming0625

brainpy/dyn/runners.py

@@ -566,8 +566,7 @@ def f_predict(self, shared_args: Dict = None):
 
       monitor_func = self.build_monitors(self._mon_info[0], self._mon_info[1], shared_args)
 
-      def _step_func(inputs):
-        t, i, x = inputs
+      def _step_func(t, i, x):
         self.target.clear_input()
         # input step
         shared = DotDict(t=t, i=i, dt=self.dt)
@@ -586,8 +585,7 @@ def _step_func(inputs):
       if self.jit['predict']:
         dyn_vars = self.target.vars()
         dyn_vars.update(self.dyn_vars)
-        f = bm.make_loop(_step_func, dyn_vars=dyn_vars.unique(), has_return=True)
-        run_func = lambda all_inputs: f(all_inputs)[1]
+        run_func = lambda all_inputs: bm.for_loop(_step_func, dyn_vars.unique(), all_inputs)
 
       else:
         def run_func(xs):
@@ -601,7 +599,7 @@ def run_func(xs):
             x = tree_map(lambda x: x[i], xs, is_leaf=lambda x: isinstance(x, bm.JaxArray))
 
             # step at the i
-            output, mon = _step_func((times[i], indices[i], x))
+            output, mon = _step_func(times[i], indices[i], x)
 
             # append output and monitor
             outputs.append(output)

brainpy/integrators/runner.py

@@ -217,16 +217,12 @@ def __init__(
 
     # build the update step
     if self.jit['predict']:
-      _loop_func = bm.make_loop(
-        self._step,
-        dyn_vars=self.dyn_vars,
-        out_vars={k: self.variables[k] for k in self.monitors.keys()},
-        has_return=True
-      )
+      def _loop_func(times):
+        return bm.for_loop(self._step, self.dyn_vars, times)
     else:
       def _loop_func(times):
-        out_vars = {k: [] for k in self.monitors.keys()}
         returns = {k: [] for k in self.fun_monitors.keys()}
+        returns.update({k: [] for k in self.monitors.keys()})
         for i in range(len(times)):
           _t = times[i]
           _dt = self.dt
@@ -237,9 +233,9 @@ def _loop_func(times):
           self._step(_t)
           # variable monitors
           for k in self.monitors.keys():
-            out_vars[k].append(bm.as_device_array(self.variables[k]))
-        out_vars = {k: bm.asarray(out_vars[k]) for k in self.monitors.keys()}
-        return out_vars, returns
+            returns[k].append(bm.as_device_array(self.variables[k]))
+        returns = {k: bm.asarray(returns[k]) for k in returns.keys()}
+        return returns
     self.step_func = _loop_func
 
   def _step(self, t):
@@ -252,11 +248,6 @@ def _step(self, t):
       kwargs.update({k: v[self.idx.value] for k, v in self._dyn_args.items()})
       self.idx += 1
 
-    # return of function monitors
-    returns = dict()
-    for key, func in self.fun_monitors.items():
-      returns[key] = func(t, self.dt)
-
     # call integrator function
     update_values = self.target(**kwargs)
     if len(self.target.variables) == 1:
@@ -268,6 +259,13 @@ def _step(self, t):
     # progress bar
     if self.progress_bar:
       id_tap(lambda *args: self._pbar.update(), ())
+
+    # return of function monitors
+    returns = dict()
+    for key, func in self.fun_monitors.items():
+      returns[key] = func(t, self.dt)
+    for k in self.monitors.keys():
+      returns[k] = self.variables[k].value
     return returns
 
   def run(self, duration, start_t=None, eval_time=False):
@@ -302,14 +300,13 @@ def run(self, duration, start_t=None, eval_time=False):
                                  refresh=True)
     if eval_time:
       t0 = time.time()
-    hists, returns = self.step_func(times)
+    hists = self.step_func(times)
     if eval_time:
       running_time = time.time() - t0
     if self.progress_bar:
       self._pbar.close()
 
     # post-running
-    hists.update(returns)
     times += self.dt
     if self.numpy_mon_after_run:
       times = np.asarray(times)

brainpy/train/online.py

@@ -234,8 +234,7 @@ def _make_fit_func(self, shared_args: Dict):
 
     monitor_func = self.build_monitors(self._mon_info[0], self._mon_info[1], shared_args)
 
-    def _step_func(all_inputs):
-      t, i, x, ys = all_inputs
+    def _step_func(t, i, x, ys):
       shared = DotDict(t=t, dt=self.dt, i=i)
 
       # input step
@@ -262,8 +261,7 @@ def _step_func(all_inputs):
     if self.jit['fit']:
       dyn_vars = self.target.vars()
       dyn_vars.update(self.dyn_vars)
-      f = bm.make_loop(_step_func, dyn_vars=dyn_vars.unique(), has_return=True)
-      return lambda all_inputs: f(all_inputs)[1]
+      return lambda all_inputs: bm.for_loop(_step_func, dyn_vars.unique(), all_inputs)
 
     else:
       def run_func(all_inputs):
@@ -273,7 +271,7 @@ def run_func(all_inputs):
         for i in range(times.shape[0]):
           x = tree_map(lambda x: x[i], xs)
           y = tree_map(lambda x: x[i], ys)
-          output, mon = _step_func((times[i], indices[i], x, y))
+          output, mon = _step_func(times[i], indices[i], x, y)
           outputs.append(output)
           for key, value in mon.items():
             monitors[key].append(value)