[ci] Fix bug in test_ndarray.py when using the latest version of JAX (

#708) * Update test_ndarray.py * Update ndarray.py
brainpy · Dec 16, 2024 · a08ad48 · a08ad48
1 parent c289edd
commit a08ad48
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/brainpy/_src/math/ndarray.py b/brainpy/_src/math/ndarray.py
@@ -98,7 +98,7 @@ def _check_tracer(self):
     self_value = self.value
     if hasattr(self_value, '_trace') and hasattr(self_value._trace.main, 'jaxpr_stack'):
       if len(self_value._trace.main.jaxpr_stack) == 0:
-        raise RuntimeError('This Array is modified during the transformation. '
+        raise jax.errors.UnexpectedTracerError('This Array is modified during the transformation. '
                            'BrainPy only supports transformations for Variable. '
                            'Please declare it as a Variable.') from jax.core.escaped_tracer_error(self_value, None)
     return self_value

diff --git a/brainpy/_src/math/tests/test_ndarray.py b/brainpy/_src/math/tests/test_ndarray.py
@@ -62,7 +62,7 @@ def _f(self, b):
 
   def test_tracing(self):
     print(self.f(1.))
-    with self.assertRaises(RuntimeError):
+    with self.assertRaises(jax.errors.UnexpectedTracerError):
       print(self.f(bm.ones(10)))