updates

Author: chaoming0625

brainpy/_src/dyn/channels/hyperpolarization_activated.py

@@ -89,7 +89,7 @@ def derivative(self, p, t, V):
 
   def reset_state(self, V, batch_size=None):
     self.p.value = self.f_p_inf(V)
-    if batch_size is not None:
+    if isinstance(batch_size, int):
       assert self.p.shape[0] == batch_size
 
   def update(self, V):
@@ -237,7 +237,7 @@ def reset_state(self, V, C_Ca, E_Ca, batch_size=None):
     beta = (1 - inf) / tau
     self.O.value = alpha / (alpha + alpha * self.k3 * self.P1 / self.k4 + beta)
     self.OL.value = self.k3 * self.P1 * self.O / self.k4
-    if batch_size is not None:
+    if isinstance(batch_size, int):
       assert self.P1.shape[0] == batch_size
       assert self.O.shape[0] == batch_size
       assert self.OL.shape[0] == batch_size

brainpy/_src/dyn/synapses/gap_junction.py renamed to brainpy/_src/dynold/synapses/gap_junction.py

@@ -117,6 +117,20 @@ def test_for_loop_progress_bar(self):
     ys = bm.for_loop(lambda a: a, xs, progress_bar=True)
     self.assertTrue(bm.allclose(xs, ys))
 
+  def test_for_loop2(self):
+    class MyClass(bp.DynamicalSystem):
+      def __init__(self):
+        super().__init__()
+        self.a = bm.Variable(bm.zeros(1))
+
+      def update(self):
+        self.a += 1
+
+    cls = MyClass()
+    indices = bm.arange(10)
+    bm.for_loop(cls.step_run, indices)
+    self.assertTrue(bm.allclose(cls.a, 10.))
+
 
 class TestIfElse(unittest.TestCase):
   def test1(self):