changes: remove brainpy.math.vmap

Author: chaoming0625

brainpy/analysis/highdim/slow_points.py

@@ -4,6 +4,7 @@
 import warnings
 from functools import partial
 
+from jax import vmap
 import jax.numpy
 import numpy as np
 from jax.scipy.optimize import minimize
@@ -56,15 +57,15 @@ def __init__(self, f_cell, f_type='continuous', f_loss_batch=None, verbose=True)
     if f_loss_batch is None:
       if f_type == 'discrete':
         self.f_loss = bm.jit(lambda h: bm.mean((h - f_cell(h)) ** 2))
-        self.f_loss_batch = bm.jit(lambda h: bm.mean((h - bm.vmap(f_cell, auto_infer=False)(h)) ** 2, axis=1))
+        self.f_loss_batch = bm.jit(lambda h: bm.mean((h - vmap(f_cell)(h)) ** 2, axis=1))
       if f_type == 'continuous':
         self.f_loss = bm.jit(lambda h: bm.mean(f_cell(h) ** 2))
-        self.f_loss_batch = bm.jit(lambda h: bm.mean((bm.vmap(f_cell, auto_infer=False)(h)) ** 2, axis=1))
+        self.f_loss_batch = bm.jit(lambda h: bm.mean((vmap(f_cell)(h)) ** 2, axis=1))
 
     else:
       self.f_loss_batch = f_loss_batch
       self.f_loss = bm.jit(lambda h: bm.mean(f_cell(h) ** 2))
-    self.f_jacob_batch = bm.jit(bm.vmap(bm.jacobian(f_cell)))
+    self.f_jacob_batch = bm.jit(vmap(bm.jacobian(f_cell)))
 
     # essential variables
     self._losses = None
@@ -208,7 +209,7 @@ def find_fps_with_opt_solver(self, candidates, opt_method=None):
       opt_method = lambda f, x0: minimize(f, x0, method='BFGS')
     if self.verbose:
       print(f"Optimizing to find fixed points:")
-    f_opt = bm.jit(bm.vmap(lambda x0: opt_method(self.f_loss, x0)))
+    f_opt = bm.jit(vmap(lambda x0: opt_method(self.f_loss, x0)))
     res = f_opt(bm.as_device_array(candidates))
     valid_ids = jax.numpy.where(res.success)[0]
     self._fixed_points = np.asarray(res.x[valid_ids])

brainpy/analysis/lowdim/lowdim_analyzer.py

@@ -3,6 +3,7 @@
 from functools import partial
 
 import numpy as np
+from jax import vmap
 from jax import numpy as jnp
 from jax.scipy.optimize import minimize
 
@@ -262,7 +263,7 @@ def F_fx(self):
   @property
   def F_vmap_fx(self):
     if C.F_vmap_fx not in self.analyzed_results:
-      self.analyzed_results[C.F_vmap_fx] = bm.jit(bm.vmap(self.F_fx), device=self.jit_device)
+      self.analyzed_results[C.F_vmap_fx] = bm.jit(vmap(self.F_fx), device=self.jit_device)
     return self.analyzed_results[C.F_vmap_fx]
 
   @property
@@ -289,7 +290,7 @@ def F_vmap_fp_aux(self):
       # ---
       # "X": a two-dimensional matrix: (num_batch, num_var)
       # "args": a list of one-dimensional vectors, each has the shape of (num_batch,)
-      self.analyzed_results[C.F_vmap_fp_aux] = bm.jit(bm.vmap(self.F_fixed_point_aux))
+      self.analyzed_results[C.F_vmap_fp_aux] = bm.jit(vmap(self.F_fixed_point_aux))
     return self.analyzed_results[C.F_vmap_fp_aux]
 
   @property
@@ -308,7 +309,7 @@ def F_vmap_fp_opt(self):
       # ---
       # "X": a two-dimensional matrix: (num_batch, num_var)
       # "args": a list of one-dimensional vectors, each has the shape of (num_batch,)
-      self.analyzed_results[C.F_vmap_fp_opt] = bm.jit(bm.vmap(self.F_fixed_point_opt))
+      self.analyzed_results[C.F_vmap_fp_opt] = bm.jit(vmap(self.F_fixed_point_opt))
     return self.analyzed_results[C.F_vmap_fp_opt]
 
   def _get_fixed_points(self, candidates, *args, num_seg=None, tol_aux=1e-7, loss_screen=None):
@@ -501,7 +502,7 @@ def F_y_by_x_in_fy(self):
   @property
   def F_vmap_fy(self):
     if C.F_vmap_fy not in self.analyzed_results:
-      self.analyzed_results[C.F_vmap_fy] = bm.jit(bm.vmap(self.F_fy), device=self.jit_device)
+      self.analyzed_results[C.F_vmap_fy] = bm.jit(vmap(self.F_fy), device=self.jit_device)
     return self.analyzed_results[C.F_vmap_fy]
 
   @property
@@ -663,7 +664,7 @@ def _get_fx_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
 
       if self.F_x_by_y_in_fx is not None:
         utils.output("I am evaluating fx-nullcline by F_x_by_y_in_fx ...")
-        vmap_f = bm.jit(bm.vmap(self.F_x_by_y_in_fx), device=self.jit_device)
+        vmap_f = bm.jit(vmap(self.F_x_by_y_in_fx), device=self.jit_device)
         for j, pars in enumerate(par_seg):
           if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
           mesh_values = jnp.meshgrid(*((ys,) + pars))
@@ -679,7 +680,7 @@ def _get_fx_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
 
       elif self.F_y_by_x_in_fx is not None:
         utils.output("I am evaluating fx-nullcline by F_y_by_x_in_fx ...")
-        vmap_f = bm.jit(bm.vmap(self.F_y_by_x_in_fx), device=self.jit_device)
+        vmap_f = bm.jit(vmap(self.F_y_by_x_in_fx), device=self.jit_device)
         for j, pars in enumerate(par_seg):
           if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
           mesh_values = jnp.meshgrid(*((xs,) + pars))
@@ -697,9 +698,9 @@ def _get_fx_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
         utils.output("I am evaluating fx-nullcline by optimization ...")
         # auxiliary functions
         f2 = lambda y, x, *pars: self.F_fx(x, y, *pars)
-        vmap_f2 = bm.jit(bm.vmap(f2), device=self.jit_device)
-        vmap_brentq_f2 = bm.jit(bm.vmap(utils.jax_brentq(f2)), device=self.jit_device)
-        vmap_brentq_f1 = bm.jit(bm.vmap(utils.jax_brentq(self.F_fx)), device=self.jit_device)
+        vmap_f2 = bm.jit(vmap(f2), device=self.jit_device)
+        vmap_brentq_f2 = bm.jit(vmap(utils.jax_brentq(f2)), device=self.jit_device)
+        vmap_brentq_f1 = bm.jit(vmap(utils.jax_brentq(self.F_fx)), device=self.jit_device)
 
         # num segments
         for _j, Ps in enumerate(par_seg):
@@ -756,7 +757,7 @@ def _get_fy_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
 
       if self.F_x_by_y_in_fy is not None:
         utils.output("I am evaluating fy-nullcline by F_x_by_y_in_fy ...")
-        vmap_f = bm.jit(bm.vmap(self.F_x_by_y_in_fy), device=self.jit_device)
+        vmap_f = bm.jit(vmap(self.F_x_by_y_in_fy), device=self.jit_device)
         for j, pars in enumerate(par_seg):
           if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
           mesh_values = jnp.meshgrid(*((ys,) + pars))
@@ -772,7 +773,7 @@ def _get_fy_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
 
       elif self.F_y_by_x_in_fy is not None:
         utils.output("I am evaluating fy-nullcline by F_y_by_x_in_fy ...")
-        vmap_f = bm.jit(bm.vmap(self.F_y_by_x_in_fy), device=self.jit_device)
+        vmap_f = bm.jit(vmap(self.F_y_by_x_in_fy), device=self.jit_device)
         for j, pars in enumerate(par_seg):
           if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
           mesh_values = jnp.meshgrid(*((xs,) + pars))
@@ -791,9 +792,9 @@ def _get_fy_nullcline_points(self, coords=None, tol=1e-7, num_segments=1, fp_aux
 
         # auxiliary functions
         f2 = lambda y, x, *pars: self.F_fy(x, y, *pars)
-        vmap_f2 = bm.jit(bm.vmap(f2), device=self.jit_device)
-        vmap_brentq_f2 = bm.jit(bm.vmap(utils.jax_brentq(f2)), device=self.jit_device)
-        vmap_brentq_f1 = bm.jit(bm.vmap(utils.jax_brentq(self.F_fy)), device=self.jit_device)
+        vmap_f2 = bm.jit(vmap(f2), device=self.jit_device)
+        vmap_brentq_f2 = bm.jit(vmap(utils.jax_brentq(f2)), device=self.jit_device)
+        vmap_brentq_f1 = bm.jit(vmap(utils.jax_brentq(self.F_fy)), device=self.jit_device)
 
         for j, Ps in enumerate(par_seg):
           if len(par_seg.arg_id_segments[0]) > 1: utils.output(f"{C.prefix}segment {j} ...")
@@ -841,7 +842,7 @@ def _get_fp_candidates_by_aux_rank(self, num_segments=1, num_rank=100):
     xs = self.resolutions[self.x_var].value
     ys = self.resolutions[self.y_var].value
     P = tuple(self.resolutions[p].value for p in self.target_par_names)
-    f_select = bm.jit(bm.vmap(lambda vals, ids: vals[ids], in_axes=(1, 1)))
+    f_select = bm.jit(vmap(lambda vals, ids: vals[ids], in_axes=(1, 1)))
 
     # num seguments
     if isinstance(num_segments, int):
@@ -921,10 +922,10 @@ def _get_fixed_points(self, candidates, *args, tol_aux=1e-7,
 
       if self.convert_type() == C.x_by_y:
         num_seg = len(self.resolutions[self.y_var])
-        f_vmap = bm.jit(bm.vmap(self.F_y_convert[1]))
+        f_vmap = bm.jit(vmap(self.F_y_convert[1]))
       else:
         num_seg = len(self.resolutions[self.x_var])
-        f_vmap = bm.jit(bm.vmap(self.F_x_convert[1]))
+        f_vmap = bm.jit(vmap(self.F_x_convert[1]))
       # get the signs
       signs = jnp.sign(f_vmap(candidates, *args))
       signs = signs.reshape((num_seg, -1))
@@ -954,10 +955,10 @@ def _get_fixed_points(self, candidates, *args, tol_aux=1e-7,
       # get another value
       if self.convert_type() == C.x_by_y:
         y_values = fps
-        x_values = bm.jit(bm.vmap(self.F_y_convert[0]))(y_values, *args)
+        x_values = bm.jit(vmap(self.F_y_convert[0]))(y_values, *args)
       else:
         x_values = fps
-        y_values = bm.jit(bm.vmap(self.F_x_convert[0]))(x_values, *args)
+        y_values = bm.jit(vmap(self.F_x_convert[0]))(x_values, *args)
       fps = jnp.stack([x_values, y_values]).T
       return fps, selected_ids, args
 

brainpy/analysis/lowdim/lowdim_bifurcation.py

@@ -3,6 +3,7 @@
 from functools import partial
 
 import jax.numpy as jnp
+from jax import vmap
 import numpy as np
 
 import brainpy.math as bm
@@ -42,7 +43,7 @@ def __init__(self, model, target_pars, target_vars, fixed_vars=None,
   @property
   def F_vmap_dfxdx(self):
     if C.F_vmap_dfxdx not in self.analyzed_results:
-      f = bm.jit(bm.vmap(bm.vector_grad(self.F_fx, argnums=0)), device=self.jit_device)
+      f = bm.jit(vmap(bm.vector_grad(self.F_fx, argnums=0)), device=self.jit_device)
       self.analyzed_results[C.F_vmap_dfxdx] = f
     return self.analyzed_results[C.F_vmap_dfxdx]
 
@@ -159,7 +160,7 @@ def F_vmap_jacobian(self):
     if C.F_vmap_jacobian not in self.analyzed_results:
       f1 = lambda xy, *args: jnp.array([self.F_fx(xy[0], xy[1], *args),
                                         self.F_fy(xy[0], xy[1], *args)])
-      f2 = bm.jit(bm.vmap(bm.jacobian(f1)), device=self.jit_device)
+      f2 = bm.jit(vmap(bm.jacobian(f1)), device=self.jit_device)
       self.analyzed_results[C.F_vmap_jacobian] = f2
     return self.analyzed_results[C.F_vmap_jacobian]
 

brainpy/analysis/lowdim/lowdim_phase_plane.py

@@ -2,6 +2,7 @@
 
 import jax.numpy as jnp
 import numpy as np
+from jax import vmap
 
 import brainpy.math as bm
 from brainpy import errors, math
@@ -158,7 +159,7 @@ def __init__(self,
   @property
   def F_vmap_brentq_fy(self):
     if C.F_vmap_brentq_fy not in self.analyzed_results:
-      f_opt = bm.jit(bm.vmap(utils.jax_brentq(self.F_fy)))
+      f_opt = bm.jit(vmap(utils.jax_brentq(self.F_fy)))
       self.analyzed_results[C.F_vmap_brentq_fy] = f_opt
     return self.analyzed_results[C.F_vmap_brentq_fy]
 

brainpy/analysis/utils/optimization.py

@@ -4,7 +4,7 @@
 import jax.lax
 import jax.numpy as jnp
 import numpy as np
-from jax import grad, jit
+from jax import grad, jit, vmap
 from jax.flatten_util import ravel_pytree
 
 import brainpy.math as bm
@@ -197,7 +197,7 @@ def brentq_candidates(vmap_f, *values, args=()):
 
 def brentq_roots(f, starts, ends, *vmap_args, args=()):
   in_axes = (0, 0, tuple([0] * len(vmap_args)) + tuple([None] * len(args)))
-  vmap_f_opt = bm.jit(bm.vmap(jax_brentq(f), in_axes=in_axes))
+  vmap_f_opt = bm.jit(vmap(jax_brentq(f), in_axes=in_axes))
   all_args = vmap_args + args
   if len(all_args):
     res = vmap_f_opt(starts, ends, all_args)
@@ -397,7 +397,7 @@ def roots_of_1d_by_x(f, candidates, args=()):
     return fps
   starts = candidates[candidate_ids]
   ends = candidates[candidate_ids + 1]
-  f_opt = bm.jit(bm.vmap(jax_brentq(f), in_axes=(0, 0, None)))
+  f_opt = bm.jit(vmap(jax_brentq(f), in_axes=(0, 0, None)))
   res = f_opt(starts, ends, args)
   valid_idx = jnp.where(res['status'] == ECONVERGED)[0]
   fps2 = res['root'][valid_idx]
@@ -406,7 +406,7 @@ def roots_of_1d_by_x(f, candidates, args=()):
 
 def roots_of_1d_by_xy(f, starts, ends, args):
   f = f_without_jaxarray_return(f)
-  f_opt = bm.jit(bm.vmap(jax_brentq(f)))
+  f_opt = bm.jit(vmap(jax_brentq(f)))
   res = f_opt(starts, ends, (args,))
   valid_idx = jnp.where(res['status'] == ECONVERGED)[0]
   xs = res['root'][valid_idx]

brainpy/analysis/utils/others.py

@@ -1,6 +1,7 @@
 # -*- coding: utf-8 -*-
 
 import jax.numpy as jnp
+from jax import vmap
 import numpy as np
 
 import brainpy.math as bm
@@ -76,7 +77,7 @@ def get_sign(f, xs, ys):
 
 def get_sign2(f, *xyz, args=()):
   in_axes = tuple(range(len(xyz))) + tuple([None] * len(args))
-  f = bm.jit(bm.vmap(f_without_jaxarray_return(f), in_axes=in_axes))
+  f = bm.jit(vmap(f_without_jaxarray_return(f), in_axes=in_axes))
   xyz = tuple((v.value if isinstance(v, bm.JaxArray) else v) for v in xyz)
   XYZ = jnp.meshgrid(*xyz)
   XYZ = tuple(jnp.moveaxis(v, 1, 0).flatten() for v in XYZ)

brainpy/dyn/synapses/delay_coupling.py

@@ -193,7 +193,7 @@ def update(self, _t, _dt):
       variable = getattr(self.pre, var)
 
       # delay function
-      f = bm.vmap(lambda i: delay_var(self.delay_mat[i], bm.arange(self.pre.num)))  # (pre.num,)
+      f = vmap(lambda i: delay_var(self.delay_mat[i], bm.arange(self.pre.num)))  # (pre.num,)
       delays = f(bm.arange(self.post.num))  # (post.num, pre.num)
       additive = (self.conn_mat * delays).sum(axis=1)
 

brainpy/errors.py

@@ -101,7 +101,8 @@ def __init__(self, variables=None):
     else:
       raise ValueError
 
-    msg += 'While there are changed variables which are not wrapped into "dyn_vars". Please check!'
+    # msg += 'While there are changed variables which are not wrapped into "dyn_vars". Please check!'
+    msg = 'While there are changed variables which are not wrapped into "dyn_vars". Please check!'
 
     super(JaxTracerError, self).__init__(msg)
 

brainpy/math/__init__.py

@@ -46,7 +46,7 @@
 from .autograd import *
 from .controls import *
 from .jit import *
-from .parallels import *
+# from .parallels import *
 
 # settings
 from . import setting

brainpy/math/parallels.py

@@ -36,29 +36,31 @@
 ]
 
 
-def _make_vmap(func, dyn_vars, rand_vars, in_axes, out_axes,
-               batch_idx, axis_name, reduce_func, f_name=None):
+def _make_vmap(func, nonbatched_vars, batched_vars, in_axes, out_axes,
+               batch_idx, axis_name, f_name=None):
   @functools.partial(jax.vmap, in_axes=in_axes, out_axes=out_axes, axis_name=axis_name)
-  def vmapped_func(dyn_data, rand_data, *args, **kwargs):
-    dyn_vars.assign(dyn_data)
-    rand_vars.assign(rand_data)
+  def vmapped_func(nonbatched_data, batched_data, *args, **kwargs):
+    nonbatched_vars.assign(nonbatched_data)
+    batched_vars.assign(batched_data)
     out = func(*args, **kwargs)
-    dyn_changes = dyn_vars.dict()
-    rand_changes = rand_vars.dict()
-    return out, dyn_changes, rand_changes
+    nonbatched_changes = nonbatched_vars.dict()
+    batched_changes = batched_vars.dict()
+    return nonbatched_changes, batched_changes, out
 
   def call(*args, **kwargs):
-    dyn_data = dyn_vars.dict()
     n = args[batch_idx[0]].shape[batch_idx[1]]
-    rand_data = {key: val.split_keys(n) for key, val in rand_vars.items()}
+    nonbatched_data = nonbatched_vars.dict()
+    batched_data = {key: val.split_keys(n) for key, val in batched_vars.items()}
     try:
-      out, dyn_changes, rand_changes = vmapped_func(dyn_data, rand_data, *args, **kwargs)
+      out, dyn_changes, rand_changes = vmapped_func(nonbatched_data, batched_data, *args, **kwargs)
     except UnexpectedTracerError as e:
-      dyn_vars.assign(dyn_data)
-      rand_vars.assign(rand_data)
-      raise errors.JaxTracerError(variables=dyn_vars) from e
-    for key, v in dyn_changes.items(): dyn_vars[key] = reduce_func(v)
-    for key, v in rand_changes.items(): rand_vars[key] = reduce_func(v)
+      nonbatched_vars.assign(nonbatched_data)
+      batched_vars.assign(batched_data)
+      raise errors.JaxTracerError() from e
+    # for key, v in dyn_changes.items():
+    #   dyn_vars[key] = reduce_func(v)
+    # for key, v in rand_changes.items():
+    #   rand_vars[key] = reduce_func(v)
     return out
 
   return change_func_name(name=f_name, f=call) if f_name else call
@@ -256,13 +258,12 @@ def vmap(func, dyn_vars=None, batched_vars=None,
 
       # jit function
       return _make_vmap(func=func,
-                        dyn_vars=_dyn_vars,
-                        rand_vars=_rand_vars,
+                        nonbatched_vars=_dyn_vars,
+                        batched_vars=_rand_vars,
                         in_axes=in_axes,
                         out_axes=out_axes,
                         axis_name=axis_name,
-                        batch_idx=batch_idx,
-                        reduce_func=reduce_func)
+                        batch_idx=batch_idx)
 
   else:
     raise errors.BrainPyError(f'Only support instance of {Base.__name__}, or a callable '