Merge branch 'master' into master

Author: ztqakita

brainpy/_src/math/object_transform/controls.py

@@ -906,11 +906,10 @@ def _body_fun(op):
     new_vals = body_fun(*old_vals)
     if new_vals is None:
       new_vals = old_vals
-    else:
-      if isinstance(new_vals, (list, tuple)):
-        new_vals = tuple(new_vals)
-      else:
-        new_vals = (new_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))