From 78036e9b58e782001750e0d6626859488dfee63e Mon Sep 17 00:00:00 2001
From: Sichao He <1310722434@qq.com>
Date: Tue, 14 May 2024 20:35:10 +0800
Subject: [PATCH 1/2] [fix] Replace jax.experimental.host_callback with
 jax.pure_callback (#670)

* Revert "fix issue #661 (#662)"

This reverts commit 4bd18980c0aa011c024024653405f6376bc5262a.

* Replace

* Support jax==0.4.28

This reverts commit 59fb681fddd7e8aa1f67abb0fabb118dfd12116a.

* Fix JIT bugs and Replace deprecated functions
---
 brainpy/_src/analysis/highdim/slow_points.py  |  2 +-
 brainpy/_src/dyn/rates/tests/test_nvar.py     |  2 +-
 .../_src/initialize/tests/test_decay_inits.py |  4 +-
 .../ode/tests/test_ode_method_exp_euler.py    |  4 +-
 brainpy/_src/integrators/runner.py            |  3 +-
 brainpy/_src/math/ndarray.py                  | 12 +--
 .../_src/math/object_transform/controls.py    | 13 ++-
 brainpy/_src/math/object_transform/jit.py     |  3 +-
 .../object_transform/tests/test_autograd.py   | 94 +++++++++----------
 .../math/object_transform/tests/test_base.py  |  7 ++
 .../tests/test_circular_reference.py          |  2 +-
 .../object_transform/tests/test_collector.py  |  4 +-
 .../object_transform/tests/test_controls.py   |  2 +-
 .../math/object_transform/tests/test_jit.py   |  2 +-
 brainpy/_src/math/random.py                   | 76 +++++++++------
 .../_src/optimizers/tests/test_ModifyLr.py    |  2 +-
 brainpy/_src/runners.py                       | 10 +-
 brainpy/_src/train/offline.py                 |  4 +-
 brainpy/_src/train/online.py                  |  4 +-
 brainpy/check.py                              |  9 +-
 requirements-dev.txt                          |  2 +-
 21 files changed, 144 insertions(+), 117 deletions(-)

diff --git a/brainpy/_src/analysis/highdim/slow_points.py b/brainpy/_src/analysis/highdim/slow_points.py
index ee91b55a5..9f31946f2 100644
--- a/brainpy/_src/analysis/highdim/slow_points.py
+++ b/brainpy/_src/analysis/highdim/slow_points.py
@@ -329,7 +329,7 @@ def find_fps_with_gd_method(
     """
     # optimization settings
     if optimizer is None:
-      optimizer = optim.Adam(lr=optim.ExponentialDecay(0.2, 1, 0.9999),
+      optimizer = optim.Adam(lr=optim.ExponentialDecayLR(0.2, 1, 0.9999),
                              beta1=0.9, beta2=0.999, eps=1e-8)
     else:
       if not isinstance(optimizer, optim.Optimizer):
diff --git a/brainpy/_src/dyn/rates/tests/test_nvar.py b/brainpy/_src/dyn/rates/tests/test_nvar.py
index 38b578a6c..24659815c 100644
--- a/brainpy/_src/dyn/rates/tests/test_nvar.py
+++ b/brainpy/_src/dyn/rates/tests/test_nvar.py
@@ -11,7 +11,7 @@ class Test_NVAR(parameterized.TestCase):
     def test_NVAR(self,mode):
         bm.random.seed()
         input=bm.random.randn(1,5)
-        layer=bp.dnn.NVAR(num_in=5,
+        layer=bp.dyn.NVAR(num_in=5,
                           delay=10,
                           mode=mode)
         if mode in [bm.NonBatchingMode()]:
diff --git a/brainpy/_src/initialize/tests/test_decay_inits.py b/brainpy/_src/initialize/tests/test_decay_inits.py
index bbab6d26d..22e1fa023 100644
--- a/brainpy/_src/initialize/tests/test_decay_inits.py
+++ b/brainpy/_src/initialize/tests/test_decay_inits.py
@@ -14,8 +14,8 @@
 # visualization
 def mat_visualize(matrix, cmap=None):
   if cmap is None:
-    cmap = plt.cm.get_cmap('coolwarm')
-  plt.cm.get_cmap('coolwarm')
+    cmap = plt.colormaps.get_cmap('coolwarm')
+  plt.colormaps.get_cmap('coolwarm')
   im = plt.matshow(matrix, cmap=cmap)
   plt.colorbar(mappable=im, shrink=0.8, aspect=15)
   plt.show()
diff --git a/brainpy/_src/integrators/ode/tests/test_ode_method_exp_euler.py b/brainpy/_src/integrators/ode/tests/test_ode_method_exp_euler.py
index 42ad7f487..d257454ef 100644
--- a/brainpy/_src/integrators/ode/tests/test_ode_method_exp_euler.py
+++ b/brainpy/_src/integrators/ode/tests/test_ode_method_exp_euler.py
@@ -94,8 +94,8 @@ def dV(self, V, t, h, n, Iext):
 
         return dVdt
 
-      def update(self, tdi):
-        t, dt = tdi.t, tdi.dt
+      def update(self):
+        t, dt = bp.share['t'], bp.share['dt']
         V, h, n = self.integral(self.V, self.h, self.n, t, self.input, dt=dt)
         self.spike.value = bm.logical_and(self.V < self.V_th, V >= self.V_th)
         self.V.value = V
diff --git a/brainpy/_src/integrators/runner.py b/brainpy/_src/integrators/runner.py
index 11dd42f58..dae638e15 100644
--- a/brainpy/_src/integrators/runner.py
+++ b/brainpy/_src/integrators/runner.py
@@ -9,7 +9,6 @@
 import jax.numpy as jnp
 import numpy as np
 import tqdm.auto
-from jax.experimental.host_callback import id_tap
 from jax.tree_util import tree_flatten
 
 from brainpy import math as bm
@@ -245,7 +244,7 @@ def _step_fun_integrator(self, static_args, dyn_args, t, i):
 
     # progress bar
     if self.progress_bar:
-      id_tap(lambda *args: self._pbar.update(), ())
+      jax.pure_callback(lambda *args: self._pbar.update(), ())
 
     # return of function monitors
     shared = dict(t=t + self.dt, dt=self.dt, i=i)
diff --git a/brainpy/_src/math/ndarray.py b/brainpy/_src/math/ndarray.py
index 791c8d9fe..b435415d6 100644
--- a/brainpy/_src/math/ndarray.py
+++ b/brainpy/_src/math/ndarray.py
@@ -660,7 +660,7 @@ def searchsorted(self, v, side='left', sorter=None):
     """
     return _return(self.value.searchsorted(v=_as_jax_array_(v), side=side, sorter=sorter))
 
-  def sort(self, axis=-1, kind='quicksort', order=None):
+  def sort(self, axis=-1, stable=True, order=None):
     """Sort an array in-place.
 
     Parameters
@@ -668,11 +668,8 @@ def sort(self, axis=-1, kind='quicksort', order=None):
     axis : int, optional
         Axis along which to sort. Default is -1, which means sort along the
         last axis.
-    kind : {'quicksort', 'mergesort', 'heapsort', 'stable'}
-        Sorting algorithm. The default is 'quicksort'. Note that both 'stable'
-        and 'mergesort' use timsort under the covers and, in general, the
-        actual implementation will vary with datatype. The 'mergesort' option
-        is retained for backwards compatibility.
+    stable : bool, optional
+        Whether to use a stable sorting algorithm. The default is True.
     order : str or list of str, optional
         When `a` is an array with fields defined, this argument specifies
         which fields to compare first, second, etc.  A single field can
@@ -680,7 +677,8 @@ def sort(self, axis=-1, kind='quicksort', order=None):
         but unspecified fields will still be used, in the order in which
         they come up in the dtype, to break ties.
     """
-    self.value = self.value.sort(axis=axis, kind=kind, order=order)
+    self.value = self.value.sort(axis=axis, stable=stable, order=order)
+
 
   def squeeze(self, axis=None):
     """Remove axes of length one from ``a``."""
diff --git a/brainpy/_src/math/object_transform/controls.py b/brainpy/_src/math/object_transform/controls.py
index 3edeb08e8..126ca15c2 100644
--- a/brainpy/_src/math/object_transform/controls.py
+++ b/brainpy/_src/math/object_transform/controls.py
@@ -7,7 +7,6 @@
 import jax
 import jax.numpy as jnp
 from jax.errors import UnexpectedTracerError
-from jax.experimental.host_callback import id_tap
 from jax.tree_util import tree_flatten, tree_unflatten
 from tqdm.auto import tqdm
 
@@ -421,14 +420,14 @@ def call(pred, x=None):
 def _warp(f):
   @functools.wraps(f)
   def new_f(*args, **kwargs):
-    return jax.tree_map(_as_jax_array_, f(*args, **kwargs), is_leaf=lambda a: isinstance(a, Array))
+    return jax.tree.map(_as_jax_array_, f(*args, **kwargs), is_leaf=lambda a: isinstance(a, Array))
 
   return new_f
 
 
 def _warp_data(data):
   def new_f(*args, **kwargs):
-    return jax.tree_map(_as_jax_array_, data, is_leaf=lambda a: isinstance(a, Array))
+    return jax.tree.map(_as_jax_array_, data, is_leaf=lambda a: isinstance(a, Array))
 
   return new_f
 
@@ -727,7 +726,7 @@ def fun2scan(carry, x):
       dyn_vars[k]._value = carry[k]
     results = body_fun(*x, **unroll_kwargs)
     if progress_bar:
-      id_tap(lambda *arg: bar.update(), ())
+      jax.pure_callback(lambda *arg: bar.update(), ())
     return dyn_vars.dict_data(), results
 
   if remat:
@@ -916,15 +915,15 @@ def fun2scan(carry, x):
       dyn_vars[k]._value = dyn_vars_data[k]
     carry, results = body_fun(carry, x)
     if progress_bar:
-      id_tap(lambda *arg: bar.update(), ())
-    carry = jax.tree_map(_as_jax_array_, carry, is_leaf=lambda a: isinstance(a, Array))
+      jax.pure_callback(lambda *arg: bar.update(), ())
+    carry = jax.tree.map(_as_jax_array_, carry, is_leaf=lambda a: isinstance(a, Array))
     return (dyn_vars.dict_data(), carry), results
 
   if remat:
     fun2scan = jax.checkpoint(fun2scan)
 
   def call(init, operands):
-    init = jax.tree_map(_as_jax_array_, init, is_leaf=lambda a: isinstance(a, Array))
+    init = jax.tree.map(_as_jax_array_, init, is_leaf=lambda a: isinstance(a, Array))
     return jax.lax.scan(f=fun2scan,
                         init=(dyn_vars.dict_data(), init),
                         xs=operands,
diff --git a/brainpy/_src/math/object_transform/jit.py b/brainpy/_src/math/object_transform/jit.py
index 551a0949c..764ce1ee5 100644
--- a/brainpy/_src/math/object_transform/jit.py
+++ b/brainpy/_src/math/object_transform/jit.py
@@ -491,9 +491,8 @@ def call_fun(self, *args, **kwargs):
 
   return call_fun
 
-
 def _make_transform(fun, stack):
-  @wraps(fun)
+  # @wraps(fun)
   def _transform_function(variable_data: Dict, *args, **kwargs):
     for key, v in stack.items():
       v._value = variable_data[key]
diff --git a/brainpy/_src/math/object_transform/tests/test_autograd.py b/brainpy/_src/math/object_transform/tests/test_autograd.py
index 90829d80e..1cd7c7cd9 100644
--- a/brainpy/_src/math/object_transform/tests/test_autograd.py
+++ b/brainpy/_src/math/object_transform/tests/test_autograd.py
@@ -1172,52 +1172,52 @@ def f(a, b):
 
 
 
-class TestHessian(unittest.TestCase):
-  def test_hessian5(self):
-    bm.set_mode(bm.training_mode)
-
-    class RNN(bp.DynamicalSystem):
-      def __init__(self, num_in, num_hidden):
-        super(RNN, self).__init__()
-        self.rnn = bp.dyn.RNNCell(num_in, num_hidden, train_state=True)
-        self.out = bp.dnn.Dense(num_hidden, 1)
-
-      def update(self, x):
-        return self.out(self.rnn(x))
-
-    # define the loss function
-    def lossfunc(inputs, targets):
-      runner = bp.DSTrainer(model, progress_bar=False, numpy_mon_after_run=False)
-      predicts = runner.predict(inputs)
-      loss = bp.losses.mean_squared_error(predicts, targets)
-      return loss
-
-    model = RNN(1, 2)
-    data_x = bm.random.rand(1, 1000, 1)
-    data_y = data_x + bm.random.randn(1, 1000, 1)
-
-    bp.reset_state(model, 1)
-    losshess = bm.hessian(lossfunc, grad_vars=model.train_vars())
-    hess_matrix = losshess(data_x, data_y)
-
-    weights = model.train_vars().unique()
-
-    # define the loss function
-    def loss_func_for_jax(weight_vals, inputs, targets):
-      for k, v in weight_vals.items():
-        weights[k].value = v
-      runner = bp.DSTrainer(model, progress_bar=False, numpy_mon_after_run=False)
-      predicts = runner.predict(inputs)
-      loss = bp.losses.mean_squared_error(predicts, targets)
-      return loss
-
-    bp.reset_state(model, 1)
-    jax_hessian = jax.hessian(loss_func_for_jax, argnums=0)({k: v.value for k, v in weights.items()}, data_x, data_y)
-
-    for k, v in hess_matrix.items():
-      for kk, vv in v.items():
-        self.assertTrue(bm.allclose(vv, jax_hessian[k][kk], atol=1e-4))
-
-    bm.clear_buffer_memory()
+# class TestHessian(unittest.TestCase):
+#   def test_hessian5(self):
+#     bm.set_mode(bm.training_mode)
+#
+#     class RNN(bp.DynamicalSystem):
+#       def __init__(self, num_in, num_hidden):
+#         super(RNN, self).__init__()
+#         self.rnn = bp.dyn.RNNCell(num_in, num_hidden, train_state=True)
+#         self.out = bp.dnn.Dense(num_hidden, 1)
+#
+#       def update(self, x):
+#         return self.out(self.rnn(x))
+#
+#     # define the loss function
+#     def lossfunc(inputs, targets):
+#       runner = bp.DSTrainer(model, progress_bar=False, numpy_mon_after_run=False)
+#       predicts = runner.predict(inputs)
+#       loss = bp.losses.mean_squared_error(predicts, targets)
+#       return loss
+#
+#     model = RNN(1, 2)
+#     data_x = bm.random.rand(1, 1000, 1)
+#     data_y = data_x + bm.random.randn(1, 1000, 1)
+#
+#     bp.reset_state(model, 1)
+#     losshess = bm.hessian(lossfunc, grad_vars=model.train_vars())
+#     hess_matrix = losshess(data_x, data_y)
+#
+#     weights = model.train_vars().unique()
+#
+#     # define the loss function
+#     def loss_func_for_jax(weight_vals, inputs, targets):
+#       for k, v in weight_vals.items():
+#         weights[k].value = v
+#       runner = bp.DSTrainer(model, progress_bar=False, numpy_mon_after_run=False)
+#       predicts = runner.predict(inputs)
+#       loss = bp.losses.mean_squared_error(predicts, targets)
+#       return loss
+#
+#     bp.reset_state(model, 1)
+#     jax_hessian = jax.hessian(loss_func_for_jax, argnums=0)({k: v.value for k, v in weights.items()}, data_x, data_y)
+#
+#     for k, v in hess_matrix.items():
+#       for kk, vv in v.items():
+#         self.assertTrue(bm.allclose(vv, jax_hessian[k][kk], atol=1e-4))
+#
+#     bm.clear_buffer_memory()
 
 
diff --git a/brainpy/_src/math/object_transform/tests/test_base.py b/brainpy/_src/math/object_transform/tests/test_base.py
index 4e1923e98..d2150d51d 100644
--- a/brainpy/_src/math/object_transform/tests/test_base.py
+++ b/brainpy/_src/math/object_transform/tests/test_base.py
@@ -239,10 +239,13 @@ def test1(self):
 
     tree = jax.tree.structure(hh)
     leaves = jax.tree.leaves(hh)
+    # tree = jax.tree.structure(hh)
+    # leaves = jax.tree.leaves(hh)
 
     print(tree)
     print(leaves)
     print(jax.tree.unflatten(tree, leaves))
+    # print(jax.tree.unflatten(tree, leaves))
     print()
 
 
@@ -282,12 +285,16 @@ def all_close(x, y):
         assert bm.allclose(x, y)
 
       jax.tree.map(all_close, all_states, variables, is_leaf=bm.is_bp_array)
+      # jax.tree.map(all_close, all_states, variables, is_leaf=bm.is_bp_array)
 
       random_state = jax.tree.map(bm.random.rand_like, all_states, is_leaf=bm.is_bp_array)
       jax.tree.map(not_close, random_state, variables, is_leaf=bm.is_bp_array)
+      # random_state = jax.tree.map(bm.random.rand_like, all_states, is_leaf=bm.is_bp_array)
+      # jax.tree.map(not_close, random_state, variables, is_leaf=bm.is_bp_array)
 
       obj.load_state_dict(random_state)
       jax.tree.map(all_close, random_state, variables, is_leaf=bm.is_bp_array)
+      # jax.tree.map(all_close, random_state, variables, is_leaf=bm.is_bp_array)
 
 
 
diff --git a/brainpy/_src/math/object_transform/tests/test_circular_reference.py b/brainpy/_src/math/object_transform/tests/test_circular_reference.py
index 61606d36e..8ef89dfca 100644
--- a/brainpy/_src/math/object_transform/tests/test_circular_reference.py
+++ b/brainpy/_src/math/object_transform/tests/test_circular_reference.py
@@ -65,7 +65,7 @@ def test_nodes():
   A.pre = B
   B.pre = A
 
-  net = bp.dyn.Network(A, B)
+  net = bp.Network(A, B)
   abs_nodes = net.nodes(method='absolute')
   rel_nodes = net.nodes(method='relative')
   print()
diff --git a/brainpy/_src/math/object_transform/tests/test_collector.py b/brainpy/_src/math/object_transform/tests/test_collector.py
index 9c3d5dde6..17ba00ec9 100644
--- a/brainpy/_src/math/object_transform/tests/test_collector.py
+++ b/brainpy/_src/math/object_transform/tests/test_collector.py
@@ -7,7 +7,7 @@
 import brainpy as bp
 
 
-class GABAa_without_Variable(bp.TwoEndConn):
+class GABAa_without_Variable(bp.synapses.TwoEndConn):
   def __init__(self, pre, post, conn, delay=0., g_max=0.1, E=-75.,
                alpha=12., beta=0.1, T=1.0, T_duration=1.0, **kwargs):
     super(GABAa_without_Variable, self).__init__(pre=pre, post=post, **kwargs)
@@ -192,7 +192,7 @@ def test_neu_nodes_1():
   assert len(neu.nodes(method='relative', include_self=False)) == 1
 
 
-class GABAa_with_Variable(bp.TwoEndConn):
+class GABAa_with_Variable(bp.synapses.TwoEndConn):
   def __init__(self, pre, post, conn, delay=0., g_max=0.1, E=-75.,
                alpha=12., beta=0.1, T=1.0, T_duration=1.0, **kwargs):
     super(GABAa_with_Variable, self).__init__(pre=pre, post=post, **kwargs)
diff --git a/brainpy/_src/math/object_transform/tests/test_controls.py b/brainpy/_src/math/object_transform/tests/test_controls.py
index 7a04c2488..b48f75042 100644
--- a/brainpy/_src/math/object_transform/tests/test_controls.py
+++ b/brainpy/_src/math/object_transform/tests/test_controls.py
@@ -234,7 +234,7 @@ def f1():
                        branches=[f1,
                                  lambda: 2, lambda: 3,
                                  lambda: 4, lambda: 5],
-                       dyn_vars=var_a,
+                       # dyn_vars=var_a,
                        show_code=True)
 
     self.assertTrue(f(11) == 1)
diff --git a/brainpy/_src/math/object_transform/tests/test_jit.py b/brainpy/_src/math/object_transform/tests/test_jit.py
index d52903d43..16d0301d4 100644
--- a/brainpy/_src/math/object_transform/tests/test_jit.py
+++ b/brainpy/_src/math/object_transform/tests/test_jit.py
@@ -157,7 +157,7 @@ class MyObj:
       def __init__(self):
         self.a = bm.Variable(bm.ones(2))
 
-      @bm.cls_jit(static_argnums=1)
+      @bm.cls_jit(static_argnums=0)
       def f(self, b, c):
         self.a.value *= b
         self.a.value /= c
diff --git a/brainpy/_src/math/random.py b/brainpy/_src/math/random.py
index 9ae012bc4..74190cb2a 100644
--- a/brainpy/_src/math/random.py
+++ b/brainpy/_src/math/random.py
@@ -4,13 +4,14 @@
 from collections import namedtuple
 from functools import partial
 from operator import index
-from typing import Optional, Union, Sequence
+from typing import Optional, Union, Sequence, Any
 
 import jax
 import numpy as np
 from jax import lax, jit, vmap, numpy as jnp, random as jr, core, dtypes
 from jax._src.array import ArrayImpl
-from jax.experimental.host_callback import call
+from jax._src.core import _canonicalize_dimension, _invalid_shape_error
+from jax._src.typing import Shape
 from jax.tree_util import register_pytree_node_class
 
 from brainpy.check import jit_error_checking, jit_error_checking_no_args
@@ -34,7 +35,7 @@
   'hypergeometric', 'logseries', 'multinomial', 'multivariate_normal',
   'negative_binomial', 'noncentral_chisquare', 'noncentral_f', 'power',
   'rayleigh', 'triangular', 'vonmises', 'wald', 'weibull', 'weibull_min',
-  'zipf', 'maxwell', 't', 'orthogonal', 'loggamma', 'categorical',
+  'zipf', 'maxwell', 't', 'orthogonal', 'loggamma', 'categorical', 'canonicalize_shape',
 
   # pytorch compatibility
   'rand_like', 'randint_like', 'randn_like',
@@ -438,6 +439,22 @@ def _check_py_seq(seq):
   return jnp.asarray(seq) if isinstance(seq, (tuple, list)) else seq
 
 
+def canonicalize_shape(shape: Shape, context: str = "") -> tuple[Any, ...]:
+  """Canonicalizes and checks for errors in a user-provided shape value.
+
+  Args:
+    shape: a Python value that represents a shape.
+
+  Returns:
+    A tuple of canonical dimension values.
+  """
+  try:
+    return tuple(map(_canonicalize_dimension, shape))
+  except TypeError:
+    pass
+  raise _invalid_shape_error(shape, context)
+
+
 @register_pytree_node_class
 class RandomState(Variable):
   """RandomState that track the random generator state. """
@@ -1098,7 +1115,7 @@ def weibull_min(self, a, scale=None, size: Optional[Union[int, Sequence[int]]] =
 
   def maxwell(self, size: Optional[Union[int, Sequence[int]]] = None, key: Optional[Union[int, JAX_RAND_KEY]] = None):
     key = self.split_key() if key is None else _formalize_key(key)
-    shape = core.canonicalize_shape(_size2shape(size)) + (3,)
+    shape = canonicalize_shape(_size2shape(size)) + (3,)
     norm_rvs = jr.normal(key=key, shape=shape)
     r = jnp.linalg.norm(norm_rvs, axis=-1)
     return _return(r)
@@ -1233,9 +1250,9 @@ def zipf(self, a, size: Optional[Union[int, Sequence[int]]] = None, key: Optiona
     if size is None:
       size = jnp.shape(a)
     dtype = jax.dtypes.canonicalize_dtype(jnp.int_)
-    r = call(lambda x: np.random.zipf(x, size).astype(dtype),
-             a,
-             result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda x: np.random.zipf(x, size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          a)
     return _return(r)
 
   def power(self, a, size: Optional[Union[int, Sequence[int]]] = None, key: Optional[Union[int, JAX_RAND_KEY]] = None):
@@ -1244,9 +1261,9 @@ def power(self, a, size: Optional[Union[int, Sequence[int]]] = None, key: Option
       size = jnp.shape(a)
     size = _size2shape(size)
     dtype = jax.dtypes.canonicalize_dtype(jnp.float_)
-    r = call(lambda a: np.random.power(a=a, size=size).astype(dtype),
-             a,
-             result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda a: np.random.power(a=a, size=size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          a)
     return _return(r)
 
   def f(self, dfnum, dfden, size: Optional[Union[int, Sequence[int]]] = None,
@@ -1260,11 +1277,11 @@ def f(self, dfnum, dfden, size: Optional[Union[int, Sequence[int]]] = None,
     size = _size2shape(size)
     d = {'dfnum': dfnum, 'dfden': dfden}
     dtype = jax.dtypes.canonicalize_dtype(jnp.float_)
-    r = call(lambda x: np.random.f(dfnum=x['dfnum'],
-                                   dfden=x['dfden'],
-                                   size=size).astype(dtype),
-             d,
-             result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda x: np.random.f(dfnum=x['dfnum'],
+                                                dfden=x['dfden'],
+                                                size=size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          d)
     return _return(r)
 
   def hypergeometric(self, ngood, nbad, nsample, size: Optional[Union[int, Sequence[int]]] = None,
@@ -1280,12 +1297,12 @@ def hypergeometric(self, ngood, nbad, nsample, size: Optional[Union[int, Sequenc
     size = _size2shape(size)
     dtype = jax.dtypes.canonicalize_dtype(jnp.int_)
     d = {'ngood': ngood, 'nbad': nbad, 'nsample': nsample}
-    r = call(lambda d: np.random.hypergeometric(ngood=d['ngood'],
-                                                nbad=d['nbad'],
-                                                nsample=d['nsample'],
-                                                size=size).astype(dtype),
-             d,
-             result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda d: np.random.hypergeometric(ngood=d['ngood'],
+                                                             nbad=d['nbad'],
+                                                             nsample=d['nsample'],
+                                                             size=size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          d)
     return _return(r)
 
   def logseries(self, p, size: Optional[Union[int, Sequence[int]]] = None,
@@ -1295,9 +1312,9 @@ def logseries(self, p, size: Optional[Union[int, Sequence[int]]] = None,
       size = jnp.shape(p)
     size = _size2shape(size)
     dtype = jax.dtypes.canonicalize_dtype(jnp.int_)
-    r = call(lambda p: np.random.logseries(p=p, size=size).astype(dtype),
-             p,
-             result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda p: np.random.logseries(p=p, size=size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          p)
     return _return(r)
 
   def noncentral_f(self, dfnum, dfden, nonc, size: Optional[Union[int, Sequence[int]]] = None,
@@ -1312,11 +1329,12 @@ def noncentral_f(self, dfnum, dfden, nonc, size: Optional[Union[int, Sequence[in
     size = _size2shape(size)
     d = {'dfnum': dfnum, 'dfden': dfden, 'nonc': nonc}
     dtype = jax.dtypes.canonicalize_dtype(jnp.float_)
-    r = call(lambda x: np.random.noncentral_f(dfnum=x['dfnum'],
-                                              dfden=x['dfden'],
-                                              nonc=x['nonc'],
-                                              size=size).astype(dtype),
-             d, result_shape=jax.ShapeDtypeStruct(size, dtype))
+    r = jax.pure_callback(lambda x: np.random.noncentral_f(dfnum=x['dfnum'],
+                                                           dfden=x['dfden'],
+                                                           nonc=x['nonc'],
+                                                           size=size).astype(dtype),
+                          jax.ShapeDtypeStruct(size, dtype),
+                          d)
     return _return(r)
 
   # PyTorch compatibility #
diff --git a/brainpy/_src/optimizers/tests/test_ModifyLr.py b/brainpy/_src/optimizers/tests/test_ModifyLr.py
index 01e51016e..67e1f6378 100644
--- a/brainpy/_src/optimizers/tests/test_ModifyLr.py
+++ b/brainpy/_src/optimizers/tests/test_ModifyLr.py
@@ -28,7 +28,7 @@ def train_data():
 class RNN(bp.DynamicalSystem):
   def __init__(self, num_in, num_hidden):
     super(RNN, self).__init__()
-    self.rnn = bp.dnn.RNNCell(num_in, num_hidden, train_state=True)
+    self.rnn = bp.dyn.RNNCell(num_in, num_hidden, train_state=True)
     self.out = bp.dnn.Dense(num_hidden, 1)
 
   def update(self, x):
diff --git a/brainpy/_src/runners.py b/brainpy/_src/runners.py
index 980ef9986..806096087 100644
--- a/brainpy/_src/runners.py
+++ b/brainpy/_src/runners.py
@@ -11,7 +11,6 @@
 import jax.numpy as jnp
 import numpy as np
 import tqdm.auto
-from jax.experimental.host_callback import id_tap
 from jax.tree_util import tree_map, tree_flatten
 
 from brainpy import math as bm, tools
@@ -632,12 +631,17 @@ def _step_func_predict(self, i, *x, shared_args=None):
 
     # finally
     if self.progress_bar:
-      id_tap(lambda *arg: self._pbar.update(), ())
+      jax.pure_callback(lambda: self._pbar.update(), ())
     # share.clear_shargs()
     clear_input(self.target)
 
     if self._memory_efficient:
-      id_tap(self._step_mon_on_cpu, mon)
+      mon_shape_dtype = jax.ShapeDtypeStruct(mon.shape, mon.dtype)
+      result = jax.pure_callback(
+        self._step_mon_on_cpu,
+        mon_shape_dtype,
+        mon,
+      )
       return out, None
     else:
       return out, mon
diff --git a/brainpy/_src/train/offline.py b/brainpy/_src/train/offline.py
index 2bfa419d6..e801a29e6 100644
--- a/brainpy/_src/train/offline.py
+++ b/brainpy/_src/train/offline.py
@@ -2,9 +2,9 @@
 
 from typing import Dict, Sequence, Union, Callable, Any
 
+import jax
 import numpy as np
 import tqdm.auto
-from jax.experimental.host_callback import id_tap
 
 import brainpy.math as bm
 from brainpy import tools
@@ -219,7 +219,7 @@ def _fun_train(self,
       targets = target_data[node.name]
       node.offline_fit(targets, fit_record)
       if self.progress_bar:
-        id_tap(lambda *args: self._pbar.update(), ())
+        jax.pure_callback(lambda *args: self._pbar.update(), ())
 
   def _step_func_monitor(self):
     res = dict()
diff --git a/brainpy/_src/train/online.py b/brainpy/_src/train/online.py
index d80764f26..862db8dfe 100644
--- a/brainpy/_src/train/online.py
+++ b/brainpy/_src/train/online.py
@@ -2,9 +2,9 @@
 import functools
 from typing import Dict, Sequence, Union, Callable
 
+import jax
 import numpy as np
 import tqdm.auto
-from jax.experimental.host_callback import id_tap
 from jax.tree_util import tree_map
 
 from brainpy import math as bm, tools
@@ -252,7 +252,7 @@ def _step_func_fit(self, i, xs: Sequence, ys: Dict, shared_args=None):
 
     # finally
     if self.progress_bar:
-      id_tap(lambda *arg: self._pbar.update(), ())
+      jax.pure_callback(lambda *arg: self._pbar.update(), ())
     return out, monitors
 
   def _check_interface(self):
diff --git a/brainpy/check.py b/brainpy/check.py
index fafc0551d..1f809d840 100644
--- a/brainpy/check.py
+++ b/brainpy/check.py
@@ -7,7 +7,6 @@
 import numpy as np
 import numpy as onp
 from jax import numpy as jnp
-from jax.experimental.host_callback import id_tap
 from jax.lax import cond
 
 conn = None
@@ -570,7 +569,11 @@ def is_all_objs(targets: Any, out_as: str = 'tuple'):
 
 
 def _err_jit_true_branch(err_fun, x):
-  id_tap(err_fun, x)
+  if isinstance(x, (tuple, list)):
+    x_shape_dtype = tuple(jax.ShapeDtypeStruct(arr.shape, arr.dtype) for arr in x)
+  else:
+    x_shape_dtype = jax.ShapeDtypeStruct(x.shape, x.dtype)
+  jax.pure_callback(err_fun, x_shape_dtype, x)
   return
 
 
@@ -629,6 +632,6 @@ def true_err_fun(arg, transforms):
     raise err
 
   cond(remove_vmap(as_jax(pred)),
-       lambda: id_tap(true_err_fun, None),
+       lambda: jax.pure_callback(true_err_fun, None),
        lambda: None)
 
diff --git a/requirements-dev.txt b/requirements-dev.txt
index 641f99fde..754073f44 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -6,7 +6,7 @@ matplotlib
 msgpack
 tqdm
 pathos
-taichi
+taichi==1.7.0
 numba
 braincore
 braintools

From 4d4eea5de3ca568895d4db41be73af62bd836534 Mon Sep 17 00:00:00 2001
From: Sichao He <1310722434@qq.com>
Date: Tue, 14 May 2024 20:35:46 +0800
Subject: [PATCH 2/2] [math] Update `CustomOpByNumba` to support JAX version >=
 0.4.24 (#669)

* [math] Update `CustomOpByNumba` to support JAX version >= 0.4.24

* Update dependency_check.py

* Update dependency_check.py

* Update requirements-dev.txt
---
 brainpy/_src/dependency_check.py              |  10 +-
 brainpy/_src/math/op_register/__init__.py     |   2 +-
 .../op_register/numba_approach/__init__.py    | 187 ++++++++++++++++--
 .../op_register/tests/test_numba_based.py     |  18 ++
 4 files changed, 192 insertions(+), 25 deletions(-)

diff --git a/brainpy/_src/dependency_check.py b/brainpy/_src/dependency_check.py
index 1e1060625..75c2051f9 100644
--- a/brainpy/_src/dependency_check.py
+++ b/brainpy/_src/dependency_check.py
@@ -57,9 +57,13 @@ def import_taichi(error_if_not_found=True):
 
   if taichi is None:
     return None
-  if taichi.__version__ != _minimal_taichi_version:
-    raise RuntimeError(taichi_install_info)
-  return taichi
+  taichi_version = taichi.__version__[0] * 10000 + taichi.__version__[1] * 100 + taichi.__version__[2]
+  minimal_taichi_version = _minimal_taichi_version[0] * 10000 + _minimal_taichi_version[1] * 100 + \
+                           _minimal_taichi_version[2]
+  if taichi_version >= minimal_taichi_version:
+    return taichi
+  else:
+    raise ModuleNotFoundError(taichi_install_info)
 
 
 def raise_taichi_not_found(*args, **kwargs):
diff --git a/brainpy/_src/math/op_register/__init__.py b/brainpy/_src/math/op_register/__init__.py
index 21c222c00..7e59e8c09 100644
--- a/brainpy/_src/math/op_register/__init__.py
+++ b/brainpy/_src/math/op_register/__init__.py
@@ -1,5 +1,5 @@
 from .numba_approach import (CustomOpByNumba,
-                             register_op_with_numba,
+                             register_op_with_numba_xla,
                              compile_cpu_signature_with_numba)
 from .base import XLACustomOp
 from .utils import register_general_batching
diff --git a/brainpy/_src/math/op_register/numba_approach/__init__.py b/brainpy/_src/math/op_register/numba_approach/__init__.py
index 5bbd04e0c..8d5cd3de1 100644
--- a/brainpy/_src/math/op_register/numba_approach/__init__.py
+++ b/brainpy/_src/math/op_register/numba_approach/__init__.py
@@ -1,29 +1,41 @@
 # -*- coding: utf-8 -*-
-
+import ctypes
 from functools import partial
 from typing import Callable
 from typing import Union, Sequence
 
 import jax
-from jax.interpreters import xla, batching, ad
+from jax.interpreters import xla, batching, ad, mlir
+from jax.lib import xla_client
 from jax.tree_util import tree_map
+from jaxlib.hlo_helpers import custom_call
 
 from brainpy._src.dependency_check import import_numba
 from brainpy._src.math.ndarray import Array
 from brainpy._src.math.object_transform.base import BrainPyObject
+from brainpy._src.math.op_register.utils import _shape_to_layout
 from brainpy.errors import PackageMissingError
 from .cpu_translation import _cpu_translation, compile_cpu_signature_with_numba
 
 numba = import_numba(error_if_not_found=False)
-
+if numba is not None:
+  from numba import types, carray, cfunc
 
 __all__ = [
   'CustomOpByNumba',
-  'register_op_with_numba',
+  'register_op_with_numba_xla',
   'compile_cpu_signature_with_numba',
 ]
 
 
+def _transform_to_shapedarray(a):
+  return jax.core.ShapedArray(a.shape, a.dtype)
+
+
+def convert_shapedarray_to_shapedtypestruct(shaped_array):
+  return jax.ShapeDtypeStruct(shape=shaped_array.shape, dtype=shaped_array.dtype)
+
+
 class CustomOpByNumba(BrainPyObject):
   """Creating a XLA custom call operator with Numba JIT on CPU backend.
 
@@ -61,20 +73,35 @@ def __init__(
     # abstract evaluation function
     if eval_shape is None:
       raise ValueError('Must provide "eval_shape" for abstract evaluation.')
+    self.eval_shape = eval_shape
 
     # cpu function
     cpu_func = con_compute
 
     # register OP
-    self.op = register_op_with_numba(
-      self.name,
-      cpu_func=cpu_func,
-      out_shapes=eval_shape,
-      batching_translation=batching_translation,
-      jvp_translation=jvp_translation,
-      transpose_translation=transpose_translation,
-      multiple_results=multiple_results,
-    )
+    if jax.__version__ > '0.4.23':
+      self.op_method = 'mlir'
+      self.op = register_op_with_numba_mlir(
+        self.name,
+        cpu_func=cpu_func,
+        out_shapes=eval_shape,
+        gpu_func_translation=None,
+        batching_translation=batching_translation,
+        jvp_translation=jvp_translation,
+        transpose_translation=transpose_translation,
+        multiple_results=multiple_results,
+      )
+    else:
+      self.op_method = 'xla'
+      self.op = register_op_with_numba_xla(
+        self.name,
+        cpu_func=cpu_func,
+        out_shapes=eval_shape,
+        batching_translation=batching_translation,
+        jvp_translation=jvp_translation,
+        transpose_translation=transpose_translation,
+        multiple_results=multiple_results,
+      )
 
   def __call__(self, *args, **kwargs):
     args = tree_map(lambda a: a.value if isinstance(a, Array) else a,
@@ -85,7 +112,7 @@ def __call__(self, *args, **kwargs):
     return res
 
 
-def register_op_with_numba(
+def register_op_with_numba_xla(
     op_name: str,
     cpu_func: Callable,
     out_shapes: Union[Callable, jax.core.ShapedArray, Sequence[jax.core.ShapedArray]],
@@ -132,13 +159,6 @@ def register_op_with_numba(
     A JAX Primitive object.
   """
 
-  if jax.__version__ > '0.4.23':
-    raise RuntimeError(f'{CustomOpByNumba.__name__} and {register_op_with_numba.__name__} are '
-                       f'only supported in JAX version <= 0.4.23. \n'
-                       f'However, you can use brainpy.math.XLACustomOp to create a custom op with numba syntax. '
-                       f'For more information, please refer to the documentation: '
-                       f'https://brainpy.readthedocs.io/en/latest/tutorial_advanced/operator_custom_with_taichi.html.')
-
   if numba is None:
     raise PackageMissingError.by_purpose('numba', 'custom op with numba')
 
@@ -202,3 +222,128 @@ def abs_eval_rule(*input_shapes, **info):
     ad.primitive_transposes[prim] = transpose_translation
 
   return prim
+
+
+def _numba_mlir_cpu_translation_rule(kernel, debug: bool, ctx, *ins, **kwargs):
+  # output information
+  outs = ctx.avals_out
+  output_shapes = tuple([out.shape for out in outs])
+  output_dtypes = tuple([out.dtype for out in outs])
+  output_layouts = tuple([_shape_to_layout(out.shape) for out in outs])
+  result_types = [mlir.aval_to_ir_type(out) for out in outs]
+
+  # input information
+  avals_in = ctx.avals_in
+  input_layouts = [_shape_to_layout(a.shape) for a in avals_in]
+  input_dtypes = tuple(inp.dtype for inp in avals_in)
+  input_shapes = tuple(inp.shape for inp in avals_in)
+
+  # compiling function
+  code_scope = dict(func_to_call=kernel, input_shapes=input_shapes, input_dtypes=input_dtypes,
+                    output_shapes=output_shapes, output_dtypes=output_dtypes, carray=carray)
+  args_in = [f'in{i} = carray(input_ptrs[{i}], input_shapes[{i}], dtype=input_dtypes[{i}])'
+             for i in range(len(input_shapes))]
+  if len(output_shapes) > 1:
+    args_out = [f'out{i} = carray(output_ptrs[{i}], output_shapes[{i}], dtype=output_dtypes[{i}])'
+                for i in range(len(output_shapes))]
+    sig = types.void(types.CPointer(types.voidptr), types.CPointer(types.voidptr))
+  else:
+    args_out = [f'out0 = carray(output_ptrs, output_shapes[0], dtype=output_dtypes[0])']
+    sig = types.void(types.voidptr, types.CPointer(types.voidptr))
+  args_call = [f'out{i}' for i in range(len(output_shapes))] + [f'in{i}' for i in range(len(input_shapes))]
+  code_string = '''
+  def numba_cpu_custom_call_target(output_ptrs, input_ptrs):
+      {args_out}
+      {args_in}
+      func_to_call({args_call})
+      '''.format(args_out="\n    ".join(args_out), args_in="\n    ".join(args_in), args_call=", ".join(args_call))
+
+  if debug:
+    print(code_string)
+  exec(compile(code_string.strip(), '', 'exec'), code_scope)
+  new_f = code_scope['numba_cpu_custom_call_target']
+
+  # register
+  xla_c_rule = cfunc(sig)(new_f)
+  target_name = f'numba_custom_call_{str(xla_c_rule.address)}'
+  capsule = ctypes.pythonapi.PyCapsule_New(xla_c_rule.address, b"xla._CUSTOM_CALL_TARGET", None)
+  xla_client.register_custom_call_target(target_name, capsule, "cpu")
+
+  # call
+  return custom_call(
+    call_target_name=target_name,
+    operands=ins,
+    operand_layouts=list(input_layouts),
+    result_layouts=list(output_layouts),
+    result_types=list(result_types),
+    has_side_effect=False,
+  ).results
+
+
+def register_op_with_numba_mlir(
+    op_name: str,
+    cpu_func: Callable,
+    out_shapes: Union[Callable, jax.core.ShapedArray, Sequence[jax.core.ShapedArray]],
+    gpu_func_translation: Callable = None,
+    batching_translation: Callable = None,
+    jvp_translation: Callable = None,
+    transpose_translation: Callable = None,
+    multiple_results: bool = False,
+):
+  if numba is None:
+    raise PackageMissingError.by_purpose('numba', 'custom op with numba')
+
+  if out_shapes is None:
+    raise RuntimeError('out_shapes cannot be None. It can be a `ShapedArray` or '
+                       'a sequence of `ShapedArray`. If it is a function, it takes as input the argument '
+                       'shapes and dtypes and should return correct output shapes of `ShapedArray`.')
+
+  prim = jax.core.Primitive(op_name)
+  prim.multiple_results = multiple_results
+
+  from numba.core.dispatcher import Dispatcher
+  if not isinstance(cpu_func, Dispatcher):
+    cpu_func = numba.jit(fastmath=True, nopython=True)(cpu_func)
+
+  def abs_eval_rule(*input_shapes, **info):
+    if callable(out_shapes):
+      shapes = out_shapes(*input_shapes, **info)
+    else:
+      shapes = out_shapes
+
+    if isinstance(shapes, jax.core.ShapedArray):
+      assert not multiple_results, "multiple_results is True, while the abstract evaluation returns only one data."
+    elif isinstance(shapes, (tuple, list)):
+      assert multiple_results, "multiple_results is False, while the abstract evaluation returns multiple data."
+      for elem in shapes:
+        if not isinstance(elem, jax.core.ShapedArray):
+          raise ValueError(f'Elements in "out_shapes" must be instances of '
+                           f'jax.abstract_arrays.ShapedArray, but we got '
+                           f'{type(elem)}: {elem}')
+    else:
+      raise ValueError(f'Unknown type {type(shapes)}, only '
+                       f'supports function, ShapedArray or '
+                       f'list/tuple of ShapedArray.')
+    return shapes
+
+  prim.def_abstract_eval(abs_eval_rule)
+  prim.def_impl(partial(xla.apply_primitive, prim))
+
+  def cpu_translation_rule(ctx, *ins, **kwargs):
+    return _numba_mlir_cpu_translation_rule(cpu_func, False, ctx, *ins, **kwargs)
+
+  mlir.register_lowering(prim, cpu_translation_rule, platform='cpu')
+
+  if gpu_func_translation is not None:
+    mlir.register_lowering(prim, gpu_func_translation, platform='gpu')
+
+  if batching_translation is not None:
+    jax.interpreters.batching.primitive_batchers[prim] = batching_translation
+
+  if jvp_translation is not None:
+    jax.interpreters.ad.primitive_jvps[prim] = jvp_translation
+
+  if transpose_translation is not None:
+    jax.interpreters.ad.primitive_transposes[prim] = transpose_translation
+
+  return prim
diff --git a/brainpy/_src/math/op_register/tests/test_numba_based.py b/brainpy/_src/math/op_register/tests/test_numba_based.py
index 28b80d0f4..f7adc695c 100644
--- a/brainpy/_src/math/op_register/tests/test_numba_based.py
+++ b/brainpy/_src/math/op_register/tests/test_numba_based.py
@@ -1,5 +1,6 @@
 import jax.core
 import pytest
+from jax.core import ShapedArray
 
 import brainpy.math as bm
 from brainpy._src.dependency_check import import_numba
@@ -35,3 +36,20 @@ def test_event_ELL():
   call(1000)
   call(100)
   bm.clear_buffer_memory()
+
+# CustomOpByNumba Test
+
+def eval_shape(a):
+  b = ShapedArray(a.shape, dtype=a.dtype)
+  return b
+
+@numba.njit(parallel=True)
+def con_compute(outs, ins):
+  b = outs
+  a = ins
+  b[:] = a + 1
+
+def test_CustomOpByNumba():
+  op = bm.CustomOpByNumba(eval_shape, con_compute, multiple_results=False)
+  print(op(bm.zeros(10)))
+  assert bm.allclose(op(bm.zeros(10)), bm.ones(10))
\ No newline at end of file