Merge pull request #358 from chaoming0625/master

Automatic transformations any function/object using `brainpy.math.Variable`

Author: chaoming0625

.gitignore

@@ -223,3 +223,4 @@ cython_debug/
 /examples/training_snn_models/logs/
 /examples/training_snn_models/data/
 /docs/tutorial_advanced/data/
+/my_tests/

brainpy/__init__.py

@@ -132,10 +132,9 @@
 
 # deprecated
 from brainpy._src.math.object_transform.base import (Base as Base,
-                                                     DynVarCollector,
+                                                     ArrayCollector,
                                                      Collector as Collector, )
-globals()['ArrayCollector'] = DynVarCollector
-globals()['TensorCollector'] = DynVarCollector
+globals()['TensorCollector'] = ArrayCollector
 
 train.__dict__['DSTrainer'] = DSTrainer
 train.__dict__['BPTT'] = BPTT
@@ -150,8 +149,8 @@
 base.base.__dict__['BrainPyObject'] = BrainPyObject
 base.base.__dict__['Base'] = Base
 base.collector.__dict__['Collector'] = Collector
-base.collector.__dict__['ArrayCollector'] = DynVarCollector
-base.collector.__dict__['TensorCollector'] = DynVarCollector
+base.collector.__dict__['ArrayCollector'] = ArrayCollector
+base.collector.__dict__['TensorCollector'] = ArrayCollector
 base.function.__dict__['FunAsObject'] = math.FunAsObject
 base.function.__dict__['Function'] = math.FunAsObject
 base.io.__dict__['save_as_h5'] = checkpoints.io.save_as_h5
@@ -162,14 +161,12 @@
 base.io.__dict__['load_by_npz'] = checkpoints.io.load_by_npz
 base.io.__dict__['load_by_pkl'] = checkpoints.io.load_by_pkl
 base.io.__dict__['load_by_mat'] = checkpoints.io.load_by_mat
-base.naming.__dict__['check_name_uniqueness'] = tools.check_name_uniqueness
-base.naming.__dict__['clear_name_cache'] = tools.clear_name_cache
-base.naming.__dict__['get_unique_name'] = tools.get_unique_name
+base.naming.__dict__['clear_name_cache'] = math.clear_name_cache
 base.__dict__['BrainPyObject'] = BrainPyObject
 base.__dict__['Base'] = Base
 base.__dict__['Collector'] = Collector
-base.__dict__['ArrayCollector'] = DynVarCollector
-base.__dict__['TensorCollector'] = DynVarCollector
+base.__dict__['ArrayCollector'] = ArrayCollector
+base.__dict__['TensorCollector'] = ArrayCollector
 base.__dict__['FunAsObject'] = math.FunAsObject
 base.__dict__['Function'] = math.FunAsObject
 base.__dict__['save_as_h5'] = checkpoints.io.save_as_h5
@@ -180,9 +177,7 @@
 base.__dict__['load_by_npz'] = checkpoints.io.load_by_npz
 base.__dict__['load_by_pkl'] = checkpoints.io.load_by_pkl
 base.__dict__['load_by_mat'] = checkpoints.io.load_by_mat
-base.__dict__['check_name_uniqueness'] = tools.check_name_uniqueness
-base.__dict__['clear_name_cache'] = tools.clear_name_cache
-base.__dict__['get_unique_name'] = tools.get_unique_name
+base.__dict__['clear_name_cache'] = math.clear_name_cache
 
 
 from . import modes
@@ -252,4 +247,5 @@
 
 from brainpy._src import checking
 tools.__dict__['checking'] = checking
+tools.__dict__['clear_name_cache'] = math.clear_name_cache
 del checking

brainpy/_src/analysis/highdim/slow_points.py

@@ -6,12 +6,12 @@
 
 import jax.numpy as jnp
 import numpy as np
-from jax import vmap
+import jax
 from jax.scipy.optimize import minimize
 from jax.tree_util import tree_flatten, tree_map
 
 import brainpy._src.math as bm
-from brainpy import optimizers as optim, losses
+from brainpy import optim, losses
 from brainpy._src.analysis import utils, base, constants
 from brainpy._src.dyn.base import DynamicalSystem
 from brainpy._src.dyn.runners import check_and_format_inputs, _f_ops
@@ -132,11 +132,11 @@ def __init__(
 
     # update function
     if target_vars is None:
-      self.target_vars = bm.DynVarCollector()
+      self.target_vars = bm.ArrayCollector()
     else:
       if not isinstance(target_vars, dict):
         raise TypeError(f'"target_vars" must be a dict but we got {type(target_vars)}')
-      self.target_vars = bm.DynVarCollector(target_vars)
+      self.target_vars = bm.ArrayCollector(target_vars)
     excluded_vars = () if excluded_vars is None else excluded_vars
     if isinstance(excluded_vars, dict):
       excluded_vars = tuple(excluded_vars.values())
@@ -337,7 +337,7 @@ def find_fps_with_gd_method(
     f_eval_loss = self._get_f_eval_loss()
 
     def f_loss():
-      return f_eval_loss(tree_map(lambda a: bm.as_device_array(a),
+      return f_eval_loss(tree_map(lambda a: bm.as_jax(a),
                                   fixed_points,
                                   is_leaf=lambda x: isinstance(x, bm.Array))).mean()
 
@@ -383,10 +383,10 @@ def batch_train(start_i, n_batch):
                 f'is below tolerance {tolerance:0.10f}.')
 
     self._opt_losses = jnp.concatenate(opt_losses)
-    self._losses = f_eval_loss(tree_map(lambda a: bm.as_device_array(a),
+    self._losses = f_eval_loss(tree_map(lambda a: bm.as_jax(a),
                                         fixed_points,
                                         is_leaf=lambda x: isinstance(x, bm.Array)))
-    self._fixed_points = tree_map(lambda a: bm.as_device_array(a),
+    self._fixed_points = tree_map(lambda a: bm.as_jax(a),
                                   fixed_points,
                                   is_leaf=lambda x: isinstance(x, bm.Array))
     self._selected_ids = jnp.arange(num_candidate)
@@ -424,9 +424,7 @@ def find_fps_with_opt_solver(
       print(f"Optimizing with {opt_solver} to find fixed points:")
 
     # optimizing
-    res = f_opt(tree_map(lambda a: bm.as_device_array(a),
-                         candidates,
-                         is_leaf=lambda a: isinstance(a, bm.Array)))
+    res = f_opt(tree_map(lambda a: bm.as_jax(a), candidates, is_leaf=lambda a: isinstance(a, bm.Array)))
 
     # results
     valid_ids = jnp.where(res.success)[0]
@@ -666,12 +664,12 @@ def _get_f_eval_loss(self, ):
   def _generate_f_eval_loss(self):
     # evaluate losses of a batch of inputs
     if self.f_type == constants.DISCRETE:
-      f_eval_loss = lambda h: self.f_loss(h, vmap(self.f_cell)(h), axis=1)
+      f_eval_loss = lambda h: self.f_loss(h, jax.vmap(self.f_cell)(h), axis=1)
     else:
-      f_eval_loss = lambda h: self.f_loss(vmap(self.f_cell)(h), axis=1)
+      f_eval_loss = lambda h: self.f_loss(jax.vmap(self.f_cell)(h), axis=1)
 
     if isinstance(self.target, DynamicalSystem):
-      @bm.jit
+      @jax.jit
       def loss_func(h):
         r = f_eval_loss(h)
         for k, v in self.excluded_vars.items():
@@ -682,7 +680,7 @@ def loss_func(h):
 
       return loss_func
     else:
-      return bm.jit(f_eval_loss)
+      return jax.jit(f_eval_loss)
 
   def _get_f_for_opt_solver(self, candidates, opt_method):
     # loss function
@@ -697,17 +695,17 @@ def _get_f_for_opt_solver(self, candidates, opt_method):
 
         def f_loss(h):
           h = {key: h[indices[i]: indices[i + 1]] for i, key in enumerate(keys)}
-          return bm.as_device_array(self.f_loss(h, self.f_cell(h)))
+          return bm.as_jax(self.f_loss(h, self.f_cell(h)))
       else:
         def f_loss(h):
-          return bm.as_device_array(self.f_loss(h, self.f_cell(h)))
+          return bm.as_jax(self.f_loss(h, self.f_cell(h)))
     else:
       # overall loss function for fixed points optimization
       def f_loss(h):
         return self.f_loss(self.f_cell(h))
 
-    @bm.jit
-    @vmap
+    @jax.jit
+    @jax.vmap
     def f_opt(x0):
       for k, v in self.target_vars.items():
         v.value = x0[k] if (v.batch_axis is None) else jnp.expand_dims(x0[k], axis=v.batch_axis)
@@ -785,7 +783,7 @@ def jacob(x0):
     else:
       jacob = self.f_cell
 
-    f_jac = bm.jit(vmap(bm.jacobian(jacob)))
+    f_jac = jax.jit(jax.vmap(bm.jacobian(jacob)))
 
     if isinstance(self.target, DynamicalSystem):
       def jacobian_func(x):

brainpy/_src/analysis/utils/model.py

@@ -1,9 +1,9 @@
 # -*- coding: utf-8 -*-
 
 
-from brainpy._src.math.ndarray import Variable
+from brainpy._src.math.object_transform import Variable
 from brainpy._src.math.environment import get_float
-from brainpy._src.math.arrayinterporate import as_jax
+from brainpy._src.math.interoperability import as_jax
 from brainpy._src.dyn.base import DynamicalSystem
 from brainpy._src.dyn.runners import DSRunner
 from brainpy._src.integrators.base import Integrator

brainpy/_src/base/naming.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
 
-from brainpy._src.tools import naming
+from brainpy._src.math.object_transform import naming
 
 __all__ = [
   'check_name_uniqueness',

brainpy/_src/checkpoints/io.py

@@ -8,7 +8,7 @@
 
 from brainpy import errors
 import brainpy.math as bm
-from brainpy._src.math.object_transform.base import BrainPyObject, DynVarCollector
+from brainpy._src.math.object_transform.base import BrainPyObject, ArrayCollector
 
 
 logger = logging.getLogger('brainpy.brainpy_object.io')
@@ -120,7 +120,7 @@ def _load(
 
 
 def _unique_and_duplicate(collector: dict):
-  gather = DynVarCollector()
+  gather = ArrayCollector()
   id2name = dict()
   duplicates = ([], [])
   for k, v in collector.items():

brainpy/_src/dyn/base.py

@@ -14,7 +14,7 @@
 from brainpy._src.connect import TwoEndConnector, MatConn, IJConn, One2One, All2All
 from brainpy._src.initialize import Initializer, parameter, variable, Uniform, noise as init_noise
 from brainpy._src.integrators import odeint, sdeint
-from brainpy._src.math.ndarray import Variable, VariableView
+from brainpy._src.math.object_transform.variables import Variable, VariableView
 from brainpy._src.math.object_transform.base import BrainPyObject, Collector
 from brainpy.errors import NoImplementationError, UnsupportedError
 from brainpy.types import ArrayType, Shape

brainpy/_src/dyn/neurons/biological_models.py

@@ -317,7 +317,7 @@ def update(self, x=None):
       x = self.input.value
     else:
       x = 0. if x is None else x
-    V, m, h, n = self.integral(self.V, self.m, self.h, self.n, t, x, dt)
+    V, m, h, n = self.integral(self.V.value, self.m.value, self.h.value, self.n.value, t, x, dt)
     self.spike.value = bm.logical_and(self.V < self.V_th, V >= self.V_th)
     self.V.value = V
     self.m.value = m

brainpy/_src/dyn/runners.py

@@ -301,6 +301,13 @@ class DSRunner(Runner):
     In order to be compatible with previous API, default is set to be ``False``.
 
     .. versionadded:: 2.3.1
+
+  memory_efficient: bool
+    Whether using the memory-efficient way to just-in-time compile the given target.
+    Default is False.
+
+    .. versionadded:: 2.3.8
+
   """
 
   target: DynamicalSystem
@@ -697,3 +704,4 @@ def __del__(self):
       for key in tuple(self._f_predict_compiled.keys()):
         self._f_predict_compiled.pop(key)
     super(DSRunner, self).__del__()
+

brainpy/_src/integrators/constants.py

@@ -1,7 +1,7 @@
 # -*- coding: utf-8 -*-
 
 # import brainpy.math as bm
-from brainpy._src.tools import naming
+from brainpy._src.math.object_transform import naming
 
 __all__ = [
   'DT',