From 5f33a668993864c9928d55b49765a8b466ba9138 Mon Sep 17 00:00:00 2001
From: Chaoming Wang <chao.brain@qq.com>
Date: Tue, 18 Jun 2024 17:17:19 +0800
Subject: [PATCH] Fix surrogate gradient function and numpy 2.0 compatibility
 (#679)

* fix surrogate batching

* fix numpy 2.0 compatible issue

* fix numpy 2.0 compatible issue

* updates

* fix numpy2.0 compatible issue

* Skip the operators tests for GitHub action server

* Update

* Update

* Update

* Update test_taichi_based.py

* Update test_get_weight_matrix.py

---------

Co-authored-by: He Sichao <1310722434@qq.com>
---
 .github/workflows/CI.yml                      |  6 +--
 brainpy/_src/losses/comparison.py             |  3 +-
 brainpy/_src/math/compat_numpy.py             | 35 +++----------
 .../_src/math/event/tests/test_event_csrmm.py | 14 ++++-
 .../_src/math/event/tests/test_event_csrmv.py |  7 ++-
 .../math/jitconn/tests/test_event_matvec.py   |  5 ++
 .../jitconn/tests/test_get_weight_matrix.py   | 12 ++++-
 .../_src/math/jitconn/tests/test_matvec.py    |  5 ++
 .../op_register/tests/test_taichi_based.py    |  5 ++
 brainpy/_src/math/sparse/tests/test_csrmm.py  | 16 +++++-
 brainpy/_src/math/sparse/tests/test_csrmv.py  |  6 ++-
 brainpy/_src/math/surrogate/__init__.py       |  1 -
 brainpy/_src/math/surrogate/_one_input.py     | 11 +++-
 brainpy/_src/math/surrogate/_one_input_new.py |  3 +-
 brainpy/_src/math/surrogate/base.py           | 19 -------
 brainpy/math/compat_numpy.py                  |  6 ---
 docs/quickstart/analysis.ipynb                | 51 ++++++++++++-------
 docs/tutorial_math/control_flows.ipynb        |  4 +-
 docs/tutorial_toolbox/synaptic_weights.ipynb  |  3 +-
 19 files changed, 124 insertions(+), 88 deletions(-)
 delete mode 100644 brainpy/_src/math/surrogate/base.py

diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml
index 01bdd87f1..4dc23b352 100644
--- a/.github/workflows/CI.yml
+++ b/.github/workflows/CI.yml
@@ -51,7 +51,7 @@ jobs:
       - name: Test with pytest
         run: |
           cd brainpy
-          pytest _src/
+          export IS_GITHUB_ACTIONS=1 && pytest _src/
 
 
   test_macos:
@@ -82,7 +82,7 @@ jobs:
       - name: Test with pytest
         run: |
           cd brainpy
-          pytest _src/
+          export IS_GITHUB_ACTIONS=1 && pytest _src/
 
 
   test_windows:
@@ -113,4 +113,4 @@ jobs:
       - name: Test with pytest
         run: |
           cd brainpy
-          pytest _src/ -p no:faulthandler
+          set IS_GITHUB_ACTIONS=1 && pytest _src/
diff --git a/brainpy/_src/losses/comparison.py b/brainpy/_src/losses/comparison.py
index ad0c3ea35..59074eb7b 100644
--- a/brainpy/_src/losses/comparison.py
+++ b/brainpy/_src/losses/comparison.py
@@ -376,7 +376,8 @@ def update(self, input, target):
 
 
 def nll_loss(input, target, reduction: str = 'mean'):
-  r"""The negative log likelihood loss.
+  r"""
+  The negative log likelihood loss.
 
   The negative log likelihood loss. It is useful to train a classification
   problem with `C` classes.
diff --git a/brainpy/_src/math/compat_numpy.py b/brainpy/_src/math/compat_numpy.py
index 0eb391458..746cbe0e9 100644
--- a/brainpy/_src/math/compat_numpy.py
+++ b/brainpy/_src/math/compat_numpy.py
@@ -10,7 +10,6 @@
 from .interoperability import *
 from .ndarray import Array
 
-
 __all__ = [
   'full', 'full_like', 'eye', 'identity', 'diag', 'tri', 'tril', 'triu',
   'empty', 'empty_like', 'ones', 'ones_like', 'zeros', 'zeros_like',
@@ -92,9 +91,8 @@
   'ravel_multi_index', 'result_type', 'sort_complex', 'unpackbits', 'delete',
 
   # unique
-  'add_docstring', 'add_newdoc', 'add_newdoc_ufunc', 'array2string', 'asanyarray',
-  'ascontiguousarray', 'asfarray', 'asscalar', 'common_type', 'disp', 'genfromtxt',
-  'loadtxt', 'info', 'issubclass_', 'place', 'polydiv', 'put', 'putmask', 'safe_eval',
+  'asanyarray', 'ascontiguousarray', 'asfarray', 'asscalar', 'common_type', 'genfromtxt',
+  'loadtxt', 'info', 'place', 'polydiv', 'put', 'putmask', 'safe_eval',
   'savetxt', 'savez_compressed', 'show_config', 'typename', 'copyto', 'matrix', 'asmatrix', 'mat',
 
 ]
@@ -204,11 +202,12 @@ def ascontiguousarray(a, dtype=None, order=None):
   return asarray(a, dtype=dtype, order=order)
 
 
-def asfarray(a, dtype=np.float_):
+def asfarray(a, dtype=None):
   if not np.issubdtype(dtype, np.inexact):
-    dtype = np.float_
+    dtype = np.float64
   return asarray(a, dtype=dtype)
 
+
 def in1d(ar1, ar2, assume_unique: bool = False, invert: bool = False) -> Array:
   del assume_unique
   ar1_flat = ravel(ar1)
@@ -227,6 +226,7 @@ def in1d(ar1, ar2, assume_unique: bool = False, invert: bool = False) -> Array:
   else:
     return asarray((ar1_flat[:, None] == ar2_flat[None, :]).any(-1))
 
+
 # Others
 # ------
 meshgrid = _compatible_with_brainpy_array(jnp.meshgrid)
@@ -454,7 +454,6 @@ def msort(a):
 sometrue = any
 
 
-
 def shape(a):
   """
   Return the shape of an array.
@@ -648,7 +647,6 @@ def size(a, axis=None):
 finfo = jnp.finfo
 iinfo = jnp.iinfo
 
-
 can_cast = _compatible_with_brainpy_array(jnp.can_cast)
 choose = _compatible_with_brainpy_array(jnp.choose)
 copy = _compatible_with_brainpy_array(jnp.copy)
@@ -678,23 +676,6 @@ def size(a, axis=None):
 # Unique APIs
 # -----------
 
-add_docstring = np.add_docstring
-add_newdoc = np.add_newdoc
-add_newdoc_ufunc = np.add_newdoc_ufunc
-
-
-def array2string(a, max_line_width=None, precision=None,
-                 suppress_small=None, separator=' ', prefix="",
-                 style=np._NoValue, formatter=None, threshold=None,
-                 edgeitems=None, sign=None, floatmode=None, suffix="",
-                 legacy=None):
-  a = as_numpy(a)
-  return array2string(a, max_line_width=max_line_width, precision=precision,
-                      suppress_small=suppress_small, separator=separator, prefix=prefix,
-                      style=style, formatter=formatter, threshold=threshold,
-                      edgeitems=edgeitems, sign=sign, floatmode=floatmode, suffix=suffix,
-                      legacy=legacy)
-
 
 def asscalar(a):
   return a.item()
@@ -731,13 +712,9 @@ def common_type(*arrays):
     return array_type[0][precision]
 
 
-disp = np.disp
-
 genfromtxt = lambda *args, **kwargs: asarray(np.genfromtxt(*args, **kwargs))
 loadtxt = lambda *args, **kwargs: asarray(np.loadtxt(*args, **kwargs))
-
 info = np.info
-issubclass_ = np.issubclass_
 
 
 def place(arr, mask, vals):
diff --git a/brainpy/_src/math/event/tests/test_event_csrmm.py b/brainpy/_src/math/event/tests/test_event_csrmm.py
index 12a35ef34..0df7fc8ff 100644
--- a/brainpy/_src/math/event/tests/test_event_csrmm.py
+++ b/brainpy/_src/math/event/tests/test_event_csrmm.py
@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
-
+import os
 from functools import partial
 
 import jax
+import pytest
 from absl.testing import parameterized
 
 import brainpy as bp
@@ -10,6 +11,17 @@
 
 # bm.set_platform('gpu')
 
+import platform
+force_test = False  # turn on to force test on windows locally
+if platform.system() == 'Windows' and not force_test:
+  pytest.skip('skip windows', allow_module_level=True)
+
+
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
+
 seed = 1234
 
 
diff --git a/brainpy/_src/math/event/tests/test_event_csrmv.py b/brainpy/_src/math/event/tests/test_event_csrmv.py
index 181ee5520..e82bcdb71 100644
--- a/brainpy/_src/math/event/tests/test_event_csrmv.py
+++ b/brainpy/_src/math/event/tests/test_event_csrmv.py
@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
-
-
+import os
 from functools import partial
 
 import jax
@@ -19,6 +18,10 @@
 if platform.system() == 'Windows' and not force_test:
   pytest.skip('skip windows', allow_module_level=True)
 
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
 
 seed = 1234
 
diff --git a/brainpy/_src/math/jitconn/tests/test_event_matvec.py b/brainpy/_src/math/jitconn/tests/test_event_matvec.py
index dd1bafded..01e9acf88 100644
--- a/brainpy/_src/math/jitconn/tests/test_event_matvec.py
+++ b/brainpy/_src/math/jitconn/tests/test_event_matvec.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import os
 
 import jax
 import jax.numpy as jnp
@@ -16,6 +17,10 @@
 if platform.system() == 'Windows' and not force_test:
   pytest.skip('skip windows', allow_module_level=True)
 
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
 
 shapes = [(100, 200), (1000, 10)]
 
diff --git a/brainpy/_src/math/jitconn/tests/test_get_weight_matrix.py b/brainpy/_src/math/jitconn/tests/test_get_weight_matrix.py
index 9f10505aa..2b04ab617 100644
--- a/brainpy/_src/math/jitconn/tests/test_get_weight_matrix.py
+++ b/brainpy/_src/math/jitconn/tests/test_get_weight_matrix.py
@@ -1,4 +1,6 @@
 # -*- coding: utf-8 -*-
+import os
+
 import jax.numpy as jnp
 import pytest
 from absl.testing import parameterized
@@ -12,8 +14,14 @@
 import platform
 
 force_test = False  # turn on to force test on windows locally
-# if platform.system() == 'Windows' and not force_test:
-#   pytest.skip('skip windows', allow_module_level=True)
+if platform.system() == 'Windows' and not force_test:
+  pytest.skip('skip windows', allow_module_level=True)
+
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
+
 
 shapes = [
   (2, 2),
diff --git a/brainpy/_src/math/jitconn/tests/test_matvec.py b/brainpy/_src/math/jitconn/tests/test_matvec.py
index e42bd3695..7720011af 100644
--- a/brainpy/_src/math/jitconn/tests/test_matvec.py
+++ b/brainpy/_src/math/jitconn/tests/test_matvec.py
@@ -1,4 +1,5 @@
 # -*- coding: utf-8 -*-
+import os
 
 import jax
 import jax.numpy as jnp
@@ -16,6 +17,10 @@
 if platform.system() == 'Windows' and not force_test:
   pytest.skip('skip windows', allow_module_level=True)
 
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
 
 shapes = [(100, 200), (1000, 10)]
 
diff --git a/brainpy/_src/math/op_register/tests/test_taichi_based.py b/brainpy/_src/math/op_register/tests/test_taichi_based.py
index ea6dcadcf..85401d99b 100644
--- a/brainpy/_src/math/op_register/tests/test_taichi_based.py
+++ b/brainpy/_src/math/op_register/tests/test_taichi_based.py
@@ -5,6 +5,11 @@
 import brainpy.math as bm
 from brainpy._src.dependency_check import import_taichi
 
+import platform
+force_test = False  # turn on to force test on windows locally
+if platform.system() == 'Windows' and not force_test:
+  pytest.skip('skip windows', allow_module_level=True)
+  
 ti = import_taichi(error_if_not_found=False)
 if ti is None:
   pytest.skip('no taichi', allow_module_level=True)
diff --git a/brainpy/_src/math/sparse/tests/test_csrmm.py b/brainpy/_src/math/sparse/tests/test_csrmm.py
index e4346c841..f7947089b 100644
--- a/brainpy/_src/math/sparse/tests/test_csrmm.py
+++ b/brainpy/_src/math/sparse/tests/test_csrmm.py
@@ -1,8 +1,11 @@
 # -*- coding: utf-8 -*-
 
+
+import os
 from functools import partial
 
 import jax
+import pytest
 from absl.testing import parameterized
 
 import brainpy as bp
@@ -10,6 +13,16 @@
 
 # bm.set_platform('gpu')
 
+import platform
+force_test = False  # turn on to force test on windows locally
+if platform.system() == 'Windows' and not force_test:
+  pytest.skip('skip windows', allow_module_level=True)
+
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
+
 seed = 1234
 
 
@@ -133,7 +146,8 @@ def test_homo_grad(self, transpose, shape, homo_data):
                         argnums=0)
     r1 = dense_f1(homo_data)
     r2 = jax.grad(sum_op(bm.sparse.csrmm))(
-      bm.asarray([homo_data]), indices, indptr, matrix, shape=(shape[1], shape[0]) if transpose else (shape[0], shape[1]),
+      bm.asarray([homo_data]), indices, indptr, matrix,
+      shape=(shape[1], shape[0]) if transpose else (shape[0], shape[1]),
       transpose=transpose)
 
     self.assertTrue(bm.allclose(r1, r2))
diff --git a/brainpy/_src/math/sparse/tests/test_csrmv.py b/brainpy/_src/math/sparse/tests/test_csrmv.py
index acedcff12..25738db89 100644
--- a/brainpy/_src/math/sparse/tests/test_csrmv.py
+++ b/brainpy/_src/math/sparse/tests/test_csrmv.py
@@ -1,5 +1,5 @@
 # -*- coding: utf-8 -*-
-
+import os
 from functools import partial
 
 import jax
@@ -17,6 +17,10 @@
 if platform.system() == 'Windows' and not force_test:
   pytest.skip('skip windows', allow_module_level=True)
 
+# Skip the test in Github Actions
+IS_GITHUB_ACTIONS = os.getenv('IS_GITHUB_ACTIONS', '0')
+if IS_GITHUB_ACTIONS == '1':
+  pytest.skip('Skip the test in Github Actions', allow_module_level=True)
 
 seed = 1234
 
diff --git a/brainpy/_src/math/surrogate/__init__.py b/brainpy/_src/math/surrogate/__init__.py
index 199eac648..f88816d70 100644
--- a/brainpy/_src/math/surrogate/__init__.py
+++ b/brainpy/_src/math/surrogate/__init__.py
@@ -1,6 +1,5 @@
 # -*- coding: utf-8 -*-
 
 
-from .base import *
 from ._one_input_new import *
 from ._two_inputs import *
diff --git a/brainpy/_src/math/surrogate/_one_input.py b/brainpy/_src/math/surrogate/_one_input.py
index 59b4ab09b..9dcf8c756 100644
--- a/brainpy/_src/math/surrogate/_one_input.py
+++ b/brainpy/_src/math/surrogate/_one_input.py
@@ -8,7 +8,6 @@
 
 from brainpy._src.math.interoperability import as_jax
 from brainpy._src.math.ndarray import Array
-from .base import Surrogate
 
 __all__ = [
   'sigmoid',
@@ -32,6 +31,16 @@
 ]
 
 
+class Surrogate(object):
+  """The base surrograte gradient function."""
+
+  def __call__(self, *args, **kwargs):
+    raise NotImplementedError
+
+  def __repr__(self):
+    return f'{self.__class__.__name__}()'
+
+
 class _OneInpSurrogate(Surrogate):
   def __init__(self, forward_use_surrogate=False):
     self.forward_use_surrogate = forward_use_surrogate
diff --git a/brainpy/_src/math/surrogate/_one_input_new.py b/brainpy/_src/math/surrogate/_one_input_new.py
index bfffd88f5..ed9957261 100644
--- a/brainpy/_src/math/surrogate/_one_input_new.py
+++ b/brainpy/_src/math/surrogate/_one_input_new.py
@@ -12,6 +12,7 @@
 from brainpy._src.math.ndarray import Array
 
 __all__ = [
+  'Surrogate',
   'Sigmoid',
   'sigmoid',
   'PiecewiseQuadratic',
@@ -61,7 +62,7 @@ def _heaviside_imp(x, dx):
 
 
 def _heaviside_batching(args, axes):
-  return heaviside_p.bind(*args), axes
+  return heaviside_p.bind(*args), [axes[0]]
 
 
 def _heaviside_jvp(primals, tangents):
diff --git a/brainpy/_src/math/surrogate/base.py b/brainpy/_src/math/surrogate/base.py
deleted file mode 100644
index dceb58b5c..000000000
--- a/brainpy/_src/math/surrogate/base.py
+++ /dev/null
@@ -1,19 +0,0 @@
-
-
-__all__ = [
-  'Surrogate'
-]
-
-
-class Surrogate(object):
-  """The base surrograte gradient function."""
-  def __call__(self, *args, **kwargs):
-    raise NotImplementedError
-
-  def __repr__(self):
-    return f'{self.__class__.__name__}()'
-
-
-
-
-
diff --git a/brainpy/math/compat_numpy.py b/brainpy/math/compat_numpy.py
index ad6c8184f..2d068acbd 100644
--- a/brainpy/math/compat_numpy.py
+++ b/brainpy/math/compat_numpy.py
@@ -327,20 +327,14 @@
   sort_complex as sort_complex,
   unpackbits as unpackbits,
   delete as delete,
-  add_docstring as add_docstring,
-  add_newdoc as add_newdoc,
-  add_newdoc_ufunc as add_newdoc_ufunc,
-  array2string as array2string,
   asanyarray as asanyarray,
   ascontiguousarray as ascontiguousarray,
   asfarray as asfarray,
   asscalar as asscalar,
   common_type as common_type,
-  disp as disp,
   genfromtxt as genfromtxt,
   loadtxt as loadtxt,
   info as info,
-  issubclass_ as issubclass_,
   place as place,
   polydiv as polydiv,
   put as put,
diff --git a/docs/quickstart/analysis.ipynb b/docs/quickstart/analysis.ipynb
index d8b62de11..49a7d38d4 100644
--- a/docs/quickstart/analysis.ipynb
+++ b/docs/quickstart/analysis.ipynb
@@ -78,7 +78,8 @@
    ],
    "source": [
     "bp.__version__"
-   ]
+   ],
+   "id": "b9f06f10049d4a8c"
   },
   {
    "cell_type": "markdown",
@@ -108,7 +109,8 @@
     "\\tau {\\dot {V}}= - (V - V_\\mathrm{rest}) + \\Delta_T \\exp(\\frac{V - V_T}{\\Delta_T}) + RI \\\\\n",
     "\\mathrm{if}\\, \\, V > \\theta, \\quad  V \\gets  V_\\mathrm{reset}\n",
     "$$"
-   ]
+   ],
+   "id": "1ecf04d3eaf72477"
   },
   {
    "cell_type": "markdown",
@@ -355,7 +357,8 @@
    },
    "source": [
     "## Slow point analysis of a high-dimensional system"
-   ]
+   ],
+   "id": "f8c9850a60c8a2b0"
   },
   {
    "cell_type": "markdown",
@@ -366,7 +369,8 @@
     "BrainPy is also capable of performing fixed/slow point analysis of high-dimensional systems. Moreover, it can perform automatic linearization analysis around the fixed point.\n",
     "\n",
     "In the following, we use a gap junction coupled FitzHugh–Nagumo (FHN) network as an example to demonstrate how to find fixed/slow points of a high-dimensional system."
-   ]
+   ],
+   "id": "a42b23306515127c"
   },
   {
    "cell_type": "markdown",
@@ -375,7 +379,8 @@
    },
    "source": [
     "We first define the gap junction coupled FHN network as the normal ``DynamicalSystem`` class."
-   ]
+   ],
+   "id": "24471089c2a5247d"
   },
   {
    "cell_type": "code",
@@ -424,7 +429,8 @@
     "    self.V.value = self.int_V(self.V, t, self.w, self.Iext, dt)\n",
     "    self.w.value = self.int_w(self.w, t, self.V, dt)\n",
     "    self.Iext[:] = 0."
-   ]
+   ],
+   "id": "f9fa252ff8617c22"
   },
   {
    "cell_type": "markdown",
@@ -433,7 +439,8 @@
    },
    "source": [
     "Through simulation, we can easily find that this system has a limit cycle attractor, implying that an unstable fixed point exists."
-   ]
+   ],
+   "id": "158d89a270070b5d"
   },
   {
    "cell_type": "code",
@@ -485,7 +492,8 @@
     "bp.visualize.line_plot(runner.mon.ts, runner.mon.V, legend='V',\n",
     "                       plot_ids=list(range(model.num)),\n",
     "                       show=True)"
-   ]
+   ],
+   "id": "58e6a900ff39ff43"
   },
   {
    "cell_type": "markdown",
@@ -494,7 +502,8 @@
    },
    "source": [
     "Let's try to optimize the fixed points for this system. Note that we only take care of the variables ``V`` and ``w``. Different from the low-dimensional analyzer, we should provide the candidate fixed points or initial fixed points when using the high-dimensional analyzer."
-   ]
+   ],
+   "id": "22084ff723d5331"
   },
   {
    "cell_type": "code",
@@ -582,7 +591,8 @@
     "\n",
     "# remove the duplicate fixed points\n",
     "finder.keep_unique()"
-   ]
+   ],
+   "id": "90ce07cf2a0a8e3c"
   },
   {
    "cell_type": "code",
@@ -617,7 +627,8 @@
    "source": [
     "print('fixed points:', )\n",
     "finder.fixed_points"
-   ]
+   ],
+   "id": "45b81a4cdac07efa"
   },
   {
    "cell_type": "code",
@@ -651,7 +662,8 @@
    "source": [
     "print('fixed point losses:', )\n",
     "finder.losses"
-   ]
+   ],
+   "id": "adc5767da75ea8b2"
   },
   {
    "cell_type": "markdown",
@@ -660,7 +672,8 @@
    },
    "source": [
     "Let's perform the linearization analysis of the found fixed points, and visualize its decomposition results."
-   ]
+   ],
+   "id": "511b2ac23cb0c586"
   },
   {
    "cell_type": "code",
@@ -694,7 +707,8 @@
    ],
    "source": [
     "_ = finder.compute_jacobians(finder.fixed_points, plot=True)"
-   ]
+   ],
+   "id": "49e773e8700380ce"
   },
   {
    "cell_type": "markdown",
@@ -703,7 +717,8 @@
    },
    "source": [
     "This is an unstable fixed point, because one of its eigenvalues has the real part bigger than 1."
-   ]
+   ],
+   "id": "a1a7df686a3bb68f"
   },
   {
    "cell_type": "markdown",
@@ -712,7 +727,8 @@
    },
    "source": [
     "## Further reading"
-   ]
+   ],
+   "id": "b9e1a94399ffa6a1"
   },
   {
    "cell_type": "markdown",
@@ -723,7 +739,8 @@
     "- For more details about how to perform bifurcation analysis and phase plane analysis, please see the tutorial of [Low-dimensional Analyzers](../tutorial_analysis/lowdim_analysis.ipynb).\n",
     "- A good example of phase plane analysis and bifurcation analysis is the decision-making model, please see the tutorial in [Analysis of a Decision-making Model](../tutorial_analysis/decision_making_model.ipynb)\n",
     "- If you want to how to analyze the slow points (or fixed points) of your high-dimensional dynamical models, please see the tutorial of [High-dimensional Analyzers](../tutorial_analysis/highdim_analysis.ipynb)"
-   ]
+   ],
+   "id": "7a4f436e37606967"
   }
  ],
  "metadata": {
diff --git a/docs/tutorial_math/control_flows.ipynb b/docs/tutorial_math/control_flows.ipynb
index c545f55de..fa328783c 100644
--- a/docs/tutorial_math/control_flows.ipynb
+++ b/docs/tutorial_math/control_flows.ipynb
@@ -330,10 +330,10 @@
       "TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[1]..\n",
       "The error occurred while tracing the function <unknown> for eval_shape. This value became a tracer due to JAX operations on these lines:\n",
       "\n",
-      "  operation a\u001b[35m:f32[]\u001b[39m = convert_element_type[new_dtype=float32 weak_type=False] b\n",
+      "  operation a\u001B[35m:f32[]\u001B[39m = convert_element_type[new_dtype=float32 weak_type=False] b\n",
       "    from line D:\\codes\\projects\\brainpy-chaoming0625\\brainpy\\_src\\math\\ndarray.py:267:19 (__lt__)\n",
       "\n",
-      "  operation a\u001b[35m:bool[1]\u001b[39m = lt b c\n",
+      "  operation a\u001B[35m:bool[1]\u001B[39m = lt b c\n",
       "    from line D:\\codes\\projects\\brainpy-chaoming0625\\brainpy\\_src\\math\\ndarray.py:267:19 (__lt__)\n",
       "See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError\n"
      ]
diff --git a/docs/tutorial_toolbox/synaptic_weights.ipynb b/docs/tutorial_toolbox/synaptic_weights.ipynb
index 312fa831f..dfa6c5133 100644
--- a/docs/tutorial_toolbox/synaptic_weights.ipynb
+++ b/docs/tutorial_toolbox/synaptic_weights.ipynb
@@ -71,7 +71,8 @@
    ],
    "source": [
     "bp.__version__"
-   ]
+   ],
+   "id": "23c6c2fe06bbd897"
   },
   {
    "cell_type": "markdown",