[fix] bugs of control flows

Author: chaoming0625

brainpy/_src/math/object_transform/controls.py

@@ -905,9 +905,11 @@ def _body_fun(op):
       v._value = dyn_vals[k]
     new_vals = body_fun(*old_vals)
     if new_vals is None:
-      new_vals = tuple()
+      new_vals = old_vals
     if not isinstance(new_vals, tuple):
       new_vals = (new_vals,)
+    if isinstance(new_vals, list):
+      new_vals = tuple(new_vals)
     return dyn_vars.dict_data(), new_vals
 
   def _cond_fun(op):

brainpy/_src/math/object_transform/tests/test_controls.py

@@ -205,6 +205,26 @@ def body(x, y):
     print()
     print(res)
 
+  def test3(self):
+    a = bm.Variable(bm.zeros(1))
+    b = bm.Variable(bm.ones(1))
+
+    def cond(x, y):
+      return bm.all(a.value < 6.)
+
+    def body(x, y):
+      a.value += x
+      b.value *= y
+
+    res = bm.while_loop(body, cond, operands=(1., 1.))
+    self.assertTrue(bm.allclose(a, 6.))
+    self.assertTrue(bm.allclose(b, 1.))
+    print()
+    print(res)
+    print(a)
+    print(b)
+
+
   def test2(self):
     a = bm.Variable(bm.zeros(1))
     b = bm.Variable(bm.ones(1))

brainpy/_src/neurons/input_groups.py

@@ -147,7 +147,7 @@ def reset_state(self, batch_size=None):
 
   def update(self):
     self.spike.value = bm.zeros_like(self.spike)
-    bm.while_loop(self._cond_fun, self._body_fun, share.load('t'))
+    bm.while_loop(self._body_fun, self._cond_fun, share.load('t'))
     return self.spike.value
 
   # functions