improve oo-to-function transformation speed

brainpy/math/controls.py

@@ -169,9 +169,9 @@ def call(xs=None, length=None):
         turn_off_global_jit()
       except UnexpectedTracerError as e:
         turn_off_global_jit()
-        for v, d in zip(dyn_vars, init_values): v.value = d
+        for v, d in zip(dyn_vars, init_values): v._value = d
         raise errors.JaxTracerError(variables=dyn_vars) from e
-      for v, d in zip(dyn_vars, dyn_values): v.value = d
+      for v, d in zip(dyn_vars, dyn_values): v._value = d
       return tree_unflatten(tree, out_values), results
 
   else:
@@ -189,7 +189,7 @@ def call(xs):
         turn_off_global_jit()
         for v, d in zip(dyn_vars, init_values): v._value = d
         raise e
-      for v, d in zip(dyn_vars, dyn_values): v.value = d
+      for v, d in zip(dyn_vars, dyn_values): v._value = d
       return tree_unflatten(tree, out_values)
 
   return call
@@ -271,7 +271,7 @@ def call(x=None):
       turn_off_global_jit()
       for v, d in zip(dyn_vars, dyn_init): v._value = d
       raise e
-    for v, d in zip(dyn_vars, dyn_values): v.value = d
+    for v, d in zip(dyn_vars, dyn_values): v._value = d
 
   return call
 
@@ -359,7 +359,7 @@ def call(pred, x=None):
         turn_off_global_jit()
         for v, d in zip(dyn_vars, old_values): v._value = d
         raise e
-      for v, d in zip(dyn_vars, dyn_values): v.value = d
+      for v, d in zip(dyn_vars, dyn_values): v._value = d
       return res
 
   else:
@@ -477,7 +477,7 @@ def _false_fun(op):
       turn_off_global_jit()
       for v, d in zip(dyn_vars, old_values): v._value = d
       raise e
-    for v, d in zip(dyn_vars, dyn_values): v.value = d
+    for v, d in zip(dyn_vars, dyn_values): v._value = d
   else:
     turn_on_global_jit()
     res = lax.cond(pred, true_fun, false_fun, operands)
@@ -663,7 +663,7 @@ def fun2scan(dyn_vals, x):
     turn_off_global_jit()
     for v, d in zip(dyn_vars, init_vals): v._value = d
     raise e
-  for v, d in zip(dyn_vars, dyn_vals): v.value = d
+  for v, d in zip(dyn_vars, dyn_vals): v._value = d
   return out_vals
 
 
@@ -729,4 +729,4 @@ def _cond_fun(op):
     turn_off_global_jit()
     for v, d in zip(dyn_vars, dyn_init): v._value = d
     raise e
-  for v, d in zip(dyn_vars, dyn_values): v.value = d
+  for v, d in zip(dyn_vars, dyn_values): v._value = d

brainpy/math/jit.py

@@ -56,7 +56,7 @@ def call(*args, **kwargs):
       turn_off_global_jit()
       for key, v in vars.items(): v._value = variable_data[key]
       raise e
-    vars.assign(changes)
+    for key, v in vars.items(): v._value = changes[key]
     return out
 
   return change_func_name(name=f_name, f=call) if f_name else call