[oo transform] fix bugs on OO-transform changes

Author: chaoming0625

brainpy/_src/analysis/lowdim/tests/test_phase_plane.py

@@ -7,7 +7,7 @@
 import jax.numpy as jnp
 
 
-block = False
+show = False
 
 
 class TestPhasePlane(unittest.TestCase):
@@ -27,7 +27,8 @@ def int_x(x, t, Iext):
     plt.ion()
     analyzer.plot_vector_field()
     analyzer.plot_fixed_point()
-    plt.show(block=block)
+    if show:
+      plt.show()
     plt.close()
     bp.math.disable_x64()
 
@@ -74,6 +75,7 @@ def int_s2(s2, t, s1):
     analyzer.plot_vector_field()
     analyzer.plot_nullcline(coords=dict(s2='s2-s1'))
     analyzer.plot_fixed_point()
-    plt.show(block=block)
+    if show:
+      plt.show()
     plt.close()
     bp.math.disable_x64()

brainpy/_src/dyn/runners.py

@@ -668,10 +668,6 @@ def _get_f_predict(self, shared_args: Dict = None):
 
     shared_kwargs_str = serialize_kwargs(shared_args)
     if shared_kwargs_str not in self._f_predict_compiled:
-      dyn_vars = self.target.vars()
-      dyn_vars.update(self._dyn_vars)
-      dyn_vars.update(self.vars(level=0))
-      dyn_vars = dyn_vars.unique()
 
       if self._memory_efficient:
         _jit_step = bm.jit(partial(self._step_func_predict, shared_args))
@@ -688,11 +684,10 @@ def run_func(all_inputs):
           return outs, None
 
       else:
-        @bm.jit
+        step = partial(self._step_func_predict, shared_args)
+
         def run_func(all_inputs):
-          return bm.for_loop(partial(self._step_func_predict, shared_args),
-                             all_inputs,
-                             jit=self.jit['predict'])
+          return bm.for_loop(step, all_inputs, jit=self.jit['predict'])
 
       self._f_predict_compiled[shared_kwargs_str] = run_func
 

brainpy/_src/integrators/sde/normal.py

@@ -292,7 +292,7 @@ def _get_g_grad(self, f, allow_raise=False, need_grad=True):
         if not allow_raise:
           raise e
       if need_grad:
-        res[0] = bm.vector_grad(f, argnums=0, dyn_vars=self.dyn_vars)
+        res[0] = bm.vector_grad(f, argnums=0)
       return [tuple(res)], state
 
   def step(self, *args, **kwargs):

brainpy/_src/integrators/sde/tests/test_normal.py

@@ -7,7 +7,7 @@
 import matplotlib.pyplot as plt
 from brainpy._src.integrators.sde.normal import ExponentialEuler
 
-block = False
+show = False
 
 
 class TestExpEuler(unittest.TestCase):
@@ -33,16 +33,18 @@ def lorenz_g(x, y, z, t, **kwargs):
     runner.run(100.)
 
     plt.plot(runner.mon.x.flatten(), runner.mon.y.flatten())
-    plt.show(block=block)
+    if show:
+      plt.show()
+    plt.close()
 
   def test2(self):
     p = 0.1
     p2 = 0.02
 
     def lorenz_g(x, y, z, t, **kwargs):
       return bp.math.asarray([p * x, p2 * x]), \
-             bp.math.asarray([p * y, p2 * y]), \
-             bp.math.asarray([p * z, p2 * z])
+        bp.math.asarray([p * y, p2 * y]), \
+        bp.math.asarray([p * z, p2 * z])
 
     dx = lambda x, t, y, sigma=10: sigma * (y - x)
     dy = lambda y, t, x, z, rho=28: x * (rho - z) - y
@@ -54,8 +56,8 @@ def lorenz_g(x, y, z, t, **kwargs):
                             wiener_type=bp.integrators.VECTOR_WIENER,
                             var_type=bp.integrators.POP_VAR,
                             show_code=True)
-    runner = bp.integrators.IntegratorRunner(intg, monitors=['x', 'y', 'z'],
-                                             dt=0.001, inits=[1., 1., 0.], jit=False)
+    runner = bp.IntegratorRunner(intg, monitors=['x', 'y', 'z'],
+                                 dt=0.001, inits=[1., 1., 0.], jit=False)
     with self.assertRaises(ValueError):
       runner.run(100.)
 
@@ -65,8 +67,8 @@ def test3(self):
 
     def lorenz_g(x, y, z, t, **kwargs):
       return bp.math.asarray([p * x, p2 * x]).T, \
-             bp.math.asarray([p * y, p2 * y]).T, \
-             bp.math.asarray([p * z, p2 * z]).T
+        bp.math.asarray([p * y, p2 * y]).T, \
+        bp.math.asarray([p * z, p2 * z]).T
 
     dx = lambda x, t, y, sigma=10: sigma * (y - x)
     dy = lambda y, t, x, z, rho=28: x * (rho - z) - y
@@ -78,15 +80,17 @@ def lorenz_g(x, y, z, t, **kwargs):
                             wiener_type=bp.integrators.VECTOR_WIENER,
                             var_type=bp.integrators.POP_VAR,
                             show_code=True)
-    runner = bp.integrators.IntegratorRunner(intg,
-                                             monitors=['x', 'y', 'z'],
-                                             dt=0.001,
-                                             inits=[1., 1., 0.],
-                                             jit=True)
+    runner = bp.IntegratorRunner(intg,
+                                 monitors=['x', 'y', 'z'],
+                                 dt=0.001,
+                                 inits=[1., 1., 0.],
+                                 jit=True)
     runner.run(100.)
 
     plt.plot(runner.mon.x.flatten(), runner.mon.y.flatten())
-    plt.show(block=block)
+    if show:
+      plt.show()
+    plt.close()
 
 
 class TestMilstein(unittest.TestCase):
@@ -110,11 +114,14 @@ def test1(self):
                      wiener_type=bp.integrators.SCALAR_WIENER,
                      var_type=bp.integrators.POP_VAR,
                      method='milstein')
-    runner = bp.integrators.IntegratorRunner(intg,
-                                             monitors=['x', 'y', 'z'],
-                                             dt=0.001, inits=[1., 1., 0.],
-                                             jit=True)
+    runner = bp.IntegratorRunner(intg,
+                                 monitors=['x', 'y', 'z'],
+                                 dt=0.001, inits=[1., 1., 0.],
+                                 jit=True)
     runner.run(100.)
 
     plt.plot(runner.mon.x.flatten(), runner.mon.y.flatten())
-    plt.show(block=block)
+    if show:
+      plt.show()
+    plt.close()
+

brainpy/_src/math/object_transform/_tools.py

@@ -16,7 +16,9 @@ class Empty(object):
 empty = Empty()
 
 
-def _partial_fun(fun, args, kwargs,
+def _partial_fun(fun,
+                 args: tuple,
+                 kwargs: dict,
                  static_argnums: Sequence[int] = (),
                  static_argnames: Sequence[str] = ()):
   static_args, dyn_args = [], []
@@ -35,16 +37,16 @@ def _partial_fun(fun, args, kwargs,
   del args, kwargs, static_argnums, static_argnames
 
   @wraps(fun)
-  def new_fun(*dyn_args, **dyn_kwargs):
+  def new_fun(*dynargs, **dynkwargs):
     args = []
     i = 0
     for arg in static_args:
       if arg == empty:
-        args.append(dyn_args[i])
+        args.append(dynargs[i])
         i += 1
       else:
         args.append(arg)
-    return fun(*args, **static_kwargs, **dyn_kwargs)
+    return fun(*args, **static_kwargs, **dynkwargs)
 
   return new_fun, dyn_args, dyn_kwargs
 
@@ -80,14 +82,16 @@ def evaluate_dyn_vars(f,
                       static_argnames: Sequence[str] = (),
                       **kwargs):
   # TODO: better way for cache mechanism
-  if len(static_argnums) or len(static_argnames):
-    f, args, kwargs = _partial_fun(f, args, kwargs, static_argnums=static_argnums, static_argnames=static_argnames)
   stack = get_stack_cache(f)
   if stack is None:
+    if len(static_argnums) or len(static_argnames):
+      f2, args, kwargs = _partial_fun(f, args, kwargs, static_argnums=static_argnums, static_argnames=static_argnames)
+    else:
+      f2, args, kwargs = f, args, kwargs
     with jax.ensure_compile_time_eval():
       with VariableStack() as stack:
-        _ = jax.eval_shape(f, *args, *kwargs)
+        _ = jax.eval_shape(f2, *args, **kwargs)
       cache_stack(f, stack)  # cache
-      del args, kwargs
+      del args, kwargs, f2
   return stack
 

brainpy/_src/math/object_transform/jit.py

@@ -106,7 +106,8 @@ def __call__(self, *args, **kwargs):
       return self.fun(*args, **kwargs)
 
     if self._transform is None:
-      self._dyn_vars = evaluate_dyn_vars(self.fun, *args,
+      self._dyn_vars = evaluate_dyn_vars(self.fun,
+                                         *args,
                                          static_argnums=self._static_argnums,
                                          static_argnames=self._static_argnames,
                                          **kwargs)

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

@@ -100,9 +100,10 @@ def __call__(self):
     bm.random.seed(0)
 
     t = Test()
-    f_grad = bm.grad(t)
+    f_grad = bm.grad(t, grad_vars={'a': t.a, 'b': t.b, 'c': t.c})
     grads = f_grad()
-    for g in grads.values(): assert (g == 1.).all()
+    for g in grads.values():
+      assert (g == 1.).all()
 
     t = Test()
     f_grad = bm.grad(t, grad_vars=[t.a, t.b])

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

@@ -11,7 +11,7 @@
 import brainpy.math as bm
 
 
-class TestLoop(jtu.JaxTestCase):
+class TestLoop(parameterized.TestCase):
   def test_make_loop(self):
     def make_node(v1, v2):
       def update(x):