brainstate 0.1.1__py2.py3-none-any.whl → 0.1.2__py2.py3-none-any.whl

brainstate/augment/_eval_shape_test.py

@@ -14,27 +14,25 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 import unittest
 
-import brainstate as bst
+import brainstate
 
 
 class TestEvalShape(unittest.TestCase):
     def test1(self):
-        class MLP(bst.nn.Module):
+        class MLP(brainstate.nn.Module):
             def __init__(self, n_in, n_mid, n_out):
                 super().__init__()
-                self.dense1 = bst.nn.Linear(n_in, n_mid)
-                self.dense2 = bst.nn.Linear(n_mid, n_out)
+                self.dense1 = brainstate.nn.Linear(n_in, n_mid)
+                self.dense2 = brainstate.nn.Linear(n_mid, n_out)
 
             def __call__(self, x):
                 x = self.dense1(x)
-                x = bst.functional.relu(x)
+                x = brainstate.functional.relu(x)
                 x = self.dense2(x)
                 return x
 
-        r = bst.augment.abstract_init(lambda: MLP(1, 2, 3))
+        r = brainstate.augment.abstract_init(lambda: MLP(1, 2, 3))
         print(r)
-        print(bst.random.DEFAULT)
+        print(brainstate.random.DEFAULT)

brainstate/augment/_mapping_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -21,7 +20,7 @@ import jax
 import jax.numpy as jnp
 import numpy as np
 
-import brainstate as bst
+import brainstate
 import brainstate.augment
 from brainstate.augment._mapping import BatchAxisError
 from brainstate.augment._mapping import _remove_axis
@@ -29,29 +28,29 @@ from brainstate.augment._mapping import _remove_axis
 
 class TestVmap(unittest.TestCase):
     def test_vmap_1(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.State(bst.random.randn(5))
-                self.b = bst.State(bst.random.randn(5))
+                self.a = brainstate.State(brainstate.random.randn(5))
+                self.b = brainstate.State(brainstate.random.randn(5))
 
             def __call__(self, *args, **kwargs):
                 return self.a.value * self.b.value
 
         model = Model()
         r1 = model.a.value * model.b.value
-        r2 = bst.augment.vmap(model, in_states=model.states())()
+        r2 = brainstate.augment.vmap(model, in_states=model.states())()
         self.assertTrue(jnp.allclose(r1, r2))
 
     def test_vmap_2(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.ShortTermState(bst.random.randn(5))
-                self.b = bst.ShortTermState(bst.random.randn(5))
-                self.c = bst.State(bst.random.randn(1))
+                self.a = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.b = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.c = brainstate.State(brainstate.random.randn(1))
 
             def __call__(self, *args, **kwargs):
                 self.c.value = self.a.value * self.b.value
@@ -59,104 +58,104 @@ class TestVmap(unittest.TestCase):
 
         model = Model()
         with self.assertRaises(BatchAxisError):
-            r2 = bst.augment.vmap(model, in_states=model.states(bst.ShortTermState))()
+            r2 = brainstate.augment.vmap(model, in_states=model.states(brainstate.ShortTermState))()
 
         model = Model()
-        r2 = bst.augment.vmap(model, in_states=model.states(bst.ShortTermState), out_states=model.c)()
+        r2 = brainstate.augment.vmap(model, in_states=model.states(brainstate.ShortTermState), out_states=model.c)()
 
     def test_vmap_3(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.State(bst.random.randn(5))
-                self.b = bst.State(bst.random.randn(5))
+                self.a = brainstate.State(brainstate.random.randn(5))
+                self.b = brainstate.State(brainstate.random.randn(5))
 
             def __call__(self, *args, **kwargs):
                 return self.a.value * self.b.value
 
         model = Model()
         with self.assertRaises(BatchAxisError):
-            r2 = bst.augment.vmap(model, in_states=model.states(), out_states={1: model.states()})()
+            r2 = brainstate.augment.vmap(model, in_states=model.states(), out_states={1: model.states()})()
 
     def test_vmap_with_random(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.ShortTermState(bst.random.randn(5))
-                self.b = bst.ShortTermState(bst.random.randn(5))
-                self.c = bst.State(bst.random.randn(1))
+                self.a = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.b = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.c = brainstate.State(brainstate.random.randn(1))
 
             def __call__(self, key):
-                bst.random.set_key(key)
+                brainstate.random.set_key(key)
                 self.c.value = self.a.value * self.b.value
-                return self.c.value + bst.random.randn(1)
+                return self.c.value + brainstate.random.randn(1)
 
         model = Model()
-        r2 = bst.augment.vmap(
+        r2 = brainstate.augment.vmap(
             model,
-            in_states=model.states(bst.ShortTermState),
+            in_states=model.states(brainstate.ShortTermState),
             out_states=model.c
         )(
-            bst.random.split_key(5)
+            brainstate.random.split_key(5)
         )
-        print(bst.random.DEFAULT)
+        print(brainstate.random.DEFAULT)
 
     def test_vmap_with_random_v3(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.ShortTermState(bst.random.randn(5))
-                self.b = bst.ShortTermState(bst.random.randn(5))
-                self.c = bst.State(bst.random.randn(1))
+                self.a = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.b = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.c = brainstate.State(brainstate.random.randn(1))
 
             def __call__(self):
                 self.c.value = self.a.value * self.b.value
-                return self.c.value + bst.random.randn(1)
+                return self.c.value + brainstate.random.randn(1)
 
         model = Model()
-        r2 = bst.augment.vmap(
+        r2 = brainstate.augment.vmap(
             model,
-            in_states=model.states(bst.ShortTermState),
+            in_states=model.states(brainstate.ShortTermState),
             out_states=model.c
         )()
-        print(bst.random.DEFAULT)
+        print(brainstate.random.DEFAULT)
 
     def test_vmap_with_random_2(self):
-        class Model(bst.nn.Module):
+        class Model(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
 
-                self.a = bst.ShortTermState(bst.random.randn(5))
-                self.b = bst.ShortTermState(bst.random.randn(5))
-                self.c = bst.State(bst.random.randn(1))
-                self.rng = bst.random.RandomState(1)
+                self.a = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.b = brainstate.ShortTermState(brainstate.random.randn(5))
+                self.c = brainstate.State(brainstate.random.randn(1))
+                self.rng = brainstate.random.RandomState(1)
 
             def __call__(self, key):
                 self.rng.set_key(key)
                 self.c.value = self.a.value * self.b.value
-                return self.c.value + bst.random.randn(1)
+                return self.c.value + brainstate.random.randn(1)
 
         model = Model()
-        r2 = bst.augment.vmap(
+        r2 = brainstate.augment.vmap(
             model,
-            in_states=model.states(bst.ShortTermState),
+            in_states=model.states(brainstate.ShortTermState),
             out_states=model.c
         )(
-            bst.random.split_key(5)
+            brainstate.random.split_key(5)
         )
 
     def test_vmap_input(self):
-        model = bst.nn.Linear(2, 3)
+        model = brainstate.nn.Linear(2, 3)
         print(id(model), id(model.weight))
         model_id = id(model)
         weight_id = id(model.weight)
 
         x = jnp.ones((5, 2))
 
-        @bst.augment.vmap
+        @brainstate.augment.vmap
         def forward(x):
             self.assertTrue(id(model) == model_id)
             self.assertTrue(id(model.weight) == weight_id)
@@ -169,39 +168,39 @@ class TestVmap(unittest.TestCase):
         print(model.weight.value)
 
     def test_vmap_states_and_input_1(self):
-        gru = bst.nn.GRUCell(2, 3)
+        gru = brainstate.nn.GRUCell(2, 3)
         gru.init_state(5)
 
-        @bst.augment.vmap(in_states=gru.states(bst.HiddenState))
+        @brainstate.augment.vmap(in_states=gru.states(brainstate.HiddenState))
         def forward(x):
             return gru(x)
 
-        xs = bst.random.randn(5, 2)
+        xs = brainstate.random.randn(5, 2)
         y = forward(xs)
         self.assertTrue(y.shape == (5, 3))
 
     def test_vmap_jit(self):
-        class Foo(bst.nn.Module):
+        class Foo(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = bst.ParamState(jnp.arange(4))
-                self.b = bst.ShortTermState(jnp.arange(4))
+                self.a = brainstate.ParamState(jnp.arange(4))
+                self.b = brainstate.ShortTermState(jnp.arange(4))
 
             def __call__(self):
                 self.b.value = self.a.value * self.b.value
 
         foo = Foo()
 
-        @bst.augment.vmap(in_states=foo.states())
+        @brainstate.augment.vmap(in_states=foo.states())
         def mul():
             foo()
 
-        @bst.compile.jit
+        @brainstate.compile.jit
         def mul_jit(inp):
             mul()
             foo.a.value += inp
 
-        with bst.StateTraceStack() as trace:
+        with brainstate.StateTraceStack() as trace:
             mul_jit(1.)
 
         print(foo.a.value)
@@ -219,27 +218,27 @@ class TestVmap(unittest.TestCase):
         print(trace.get_read_states())
 
     def test_vmap_jit_2(self):
-        class Foo(bst.nn.Module):
+        class Foo(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = bst.ParamState(jnp.arange(4))
-                self.b = bst.ShortTermState(jnp.arange(4))
+                self.a = brainstate.ParamState(jnp.arange(4))
+                self.b = brainstate.ShortTermState(jnp.arange(4))
 
             def __call__(self):
                 self.b.value = self.a.value * self.b.value
 
         foo = Foo()
 
-        @bst.augment.vmap(in_states=foo.states())
+        @brainstate.augment.vmap(in_states=foo.states())
         def mul():
             foo()
 
-        @bst.compile.jit
+        @brainstate.compile.jit
         def mul_jit(inp):
             mul()
             foo.b.value += inp
 
-        with bst.StateTraceStack() as trace:
+        with brainstate.StateTraceStack() as trace:
             mul_jit(1.)
 
         print(foo.a.value)
@@ -258,9 +257,9 @@ class TestVmap(unittest.TestCase):
 
     def test_auto_rand_key_split(self):
         def f():
-            return bst.random.rand(1)
+            return brainstate.random.rand(1)
 
-        res = bst.augment.vmap(f, axis_size=10)()
+        res = brainstate.augment.vmap(f, axis_size=10)()
         self.assertTrue(jnp.all(~(res[0] == res[1:])))
 
         res2 = jax.vmap(f, axis_size=10)()
@@ -278,17 +277,17 @@ class TestVmap(unittest.TestCase):
         self.assertTrue(jnp.allclose(r, r2))
 
     def test_vmap_init(self):
-        class Foo(bst.nn.Module):
+        class Foo(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = bst.ParamState(jnp.arange(4))
-                self.b = bst.ShortTermState(jnp.arange(4))
+                self.a = brainstate.ParamState(jnp.arange(4))
+                self.b = brainstate.ShortTermState(jnp.arange(4))
 
             def init_state_v1(self, *args, **kwargs):
-                self.c = bst.State(jnp.arange(4))
+                self.c = brainstate.State(jnp.arange(4))
 
             def init_state_v2(self):
-                self.d = bst.State(self.c.value * 2.)
+                self.d = brainstate.State(self.c.value * 2.)
 
         foo = Foo()
 
@@ -318,11 +317,11 @@ class TestVmap(unittest.TestCase):
 class TestMap(unittest.TestCase):
     def test_map(self):
         for dim in [(10,), (10, 10), (10, 10, 10)]:
-            x = bst.random.rand(*dim)
-            r1 = bst.augment.map(lambda a: a + 1, x, batch_size=None)
-            r2 = bst.augment.map(lambda a: a + 1, x, batch_size=2)
-            r3 = bst.augment.map(lambda a: a + 1, x, batch_size=4)
-            r4 = bst.augment.map(lambda a: a + 1, x, batch_size=5)
+            x = brainstate.random.rand(*dim)
+            r1 = brainstate.augment.map(lambda a: a + 1, x, batch_size=None)
+            r2 = brainstate.augment.map(lambda a: a + 1, x, batch_size=2)
+            r3 = brainstate.augment.map(lambda a: a + 1, x, batch_size=4)
+            r4 = brainstate.augment.map(lambda a: a + 1, x, batch_size=5)
             true_r = x + 1
 
             self.assertTrue(jnp.allclose(r1, true_r))
@@ -406,7 +405,7 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
         foo = brainstate.nn.LIF(3)
         # Testing that axis_size of 0 raises an error.
         with self.assertRaises(ValueError):
-            @bst.augment.vmap_new_states(state_tag='new1', axis_size=0)
+            @brainstate.augment.vmap_new_states(state_tag='new1', axis_size=0)
             def faulty_init():
                 foo.init_state()
 
@@ -417,7 +416,7 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
         foo = brainstate.nn.LIF(3)
         # Testing that a negative axis_size raises an error.
         with self.assertRaises(ValueError):
-            @bst.augment.vmap_new_states(state_tag='new1', axis_size=-3)
+            @brainstate.augment.vmap_new_states(state_tag='new1', axis_size=-3)
             def faulty_init():
                 foo.init_state()
 
@@ -428,9 +427,9 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
 
         # Simulate an incompatible shapes scenario:
         # We intentionally assign a state with a different shape than expected.
-        @bst.augment.vmap_new_states(state_tag='new1', axis_size=5)
+        @brainstate.augment.vmap_new_states(state_tag='new1', axis_size=5)
         def faulty_init():
             # Modify state to produce an incompatible shape
-            foo.c = bst.State(jnp.arange(3))  # Original expected shape is (4,)
+            foo.c = brainstate.State(jnp.arange(3))  # Original expected shape is (4,)
 
         faulty_init()

brainstate/compile/_ad_checkpoint_test.py

@@ -13,34 +13,32 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import jax
 import jax.numpy as jnp
 from absl.testing import absltest
 
-import brainstate as bst
+import brainstate
 
 
 class TestRemat(absltest.TestCase):
     def test_basic_remat(self):
-        module = bst.compile.remat(bst.nn.Linear(2, 3))
+        module = brainstate.compile.remat(brainstate.nn.Linear(2, 3))
         y = module(jnp.ones((1, 2)))
         assert y.shape == (1, 3)
 
     def test_remat_with_scan(self):
-        class ScanLinear(bst.nn.Module):
+        class ScanLinear(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.linear = bst.nn.Linear(3, 3)
+                self.linear = brainstate.nn.Linear(3, 3)
 
             def __call__(self, x: jax.Array):
-                @bst.compile.remat
+                @brainstate.compile.remat
                 def fun(x: jax.Array, _):
                     x = self.linear(x)
                     return x, None
 
-                return bst.compile.scan(fun, x, None, length=10)[0]
+                return brainstate.compile.scan(fun, x, None, length=10)[0]
 
         m = ScanLinear()
 

brainstate/compile/_conditions_test.py

@@ -12,27 +12,26 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 import unittest
 
 import jax
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestCond(unittest.TestCase):
     def test1(self):
-        bst.random.seed(1)
-        bst.compile.cond(True, lambda: bst.random.random(10), lambda: bst.random.random(10))
-        bst.compile.cond(False, lambda: bst.random.random(10), lambda: bst.random.random(10))
+        brainstate.random.seed(1)
+        brainstate.compile.cond(True, lambda: brainstate.random.random(10), lambda: brainstate.random.random(10))
+        brainstate.compile.cond(False, lambda: brainstate.random.random(10), lambda: brainstate.random.random(10))
 
     def test2(self):
-        st1 = bst.State(bst.random.rand(10))
-        st2 = bst.State(bst.random.rand(2))
-        st3 = bst.State(bst.random.rand(5))
-        st4 = bst.State(bst.random.rand(2, 10))
+        st1 = brainstate.State(brainstate.random.rand(10))
+        st2 = brainstate.State(brainstate.random.rand(2))
+        st3 = brainstate.State(brainstate.random.rand(5))
+        st4 = brainstate.State(brainstate.random.rand(2, 10))
 
         def true_fun(x):
             st1.value = st2.value @ st4.value + x
@@ -40,7 +39,7 @@ class TestCond(unittest.TestCase):
         def false_fun(x):
             st3.value = (st3.value + 1.) * x
 
-        bst.compile.cond(True, true_fun, false_fun, 2.)
+        brainstate.compile.cond(True, true_fun, false_fun, 2.)
         assert not isinstance(st1.value, jax.core.Tracer)
         assert not isinstance(st2.value, jax.core.Tracer)
         assert not isinstance(st3.value, jax.core.Tracer)
@@ -66,7 +65,7 @@ class TestSwitch(unittest.TestCase):
                 return branches[2](x)
 
         def cfun(x):
-            return bst.compile.switch(x, branches, x)
+            return brainstate.compile.switch(x, branches, x)
 
         self.assertEqual(fun(-1), cfun(-1))
         self.assertEqual(fun(0), cfun(0))
@@ -90,7 +89,7 @@ class TestSwitch(unittest.TestCase):
             return branches[i](x, x)
 
         def cfun(x):
-            return bst.compile.switch(x, branches, x, x)
+            return brainstate.compile.switch(x, branches, x, x)
 
         self.assertEqual(fun(-1), cfun(-1))
         self.assertEqual(fun(0), cfun(0))
@@ -123,13 +122,13 @@ class TestSwitch(unittest.TestCase):
         branches3 = branches2 + [lambda x: jnp.sin(x) + jnp.cos(x)]  # requires one more residual slot
 
         def fun1(x, i):
-            return bst.compile.switch(i + 1, branches1, x)
+            return brainstate.compile.switch(i + 1, branches1, x)
 
         def fun2(x, i):
-            return bst.compile.switch(i + 1, branches2, x)
+            return brainstate.compile.switch(i + 1, branches2, x)
 
         def fun3(x, i):
-            return bst.compile.switch(i + 1, branches3, x)
+            return brainstate.compile.switch(i + 1, branches3, x)
 
         fwd1, bwd1 = get_conds(fun1)
         fwd2, bwd2 = get_conds(fun2)
@@ -149,7 +148,7 @@ class TestSwitch(unittest.TestCase):
 
     def testOneBranchSwitch(self):
         branch = lambda x: -x
-        f = lambda i, x: bst.compile.switch(i, [branch], x)
+        f = lambda i, x: brainstate.compile.switch(i, [branch], x)
         x = 7.
         self.assertEqual(f(-1, x), branch(x))
         self.assertEqual(f(0, x), branch(x))
@@ -167,12 +166,12 @@ class TestSwitch(unittest.TestCase):
 class TestIfElse(unittest.TestCase):
     def test1(self):
         def f(a):
-            return bst.compile.ifelse(conditions=[a < 0,
-                                                  a >= 0 and a < 2,
-                                                  a >= 2 and a < 5,
-                                                  a >= 5 and a < 10,
-                                                  a >= 10],
-                                      branches=[lambda: 1,
+            return brainstate.compile.ifelse(conditions=[a < 0,
+                                                         a >= 0 and a < 2,
+                                                         a >= 2 and a < 5,
+                                                         a >= 5 and a < 10,
+                                                         a >= 10],
+                                             branches=[lambda: 1,
                                                 lambda: 2,
                                                 lambda: 3,
                                                 lambda: 4,
@@ -184,38 +183,38 @@ class TestIfElse(unittest.TestCase):
 
     def test_vmap(self):
         def f(operands):
-            f = lambda a: bst.compile.ifelse([a > 10,
-                                              jnp.logical_and(a <= 10, a > 5),
-                                              jnp.logical_and(a <= 5, a > 2),
-                                              jnp.logical_and(a <= 2, a > 0),
-                                              a <= 0],
-                                             [lambda _: 1,
+            f = lambda a: brainstate.compile.ifelse([a > 10,
+                                                     jnp.logical_and(a <= 10, a > 5),
+                                                     jnp.logical_and(a <= 5, a > 2),
+                                                     jnp.logical_and(a <= 2, a > 0),
+                                                     a <= 0],
+                                                    [lambda _: 1,
                                               lambda _: 2,
                                               lambda _: 3,
                                               lambda _: 4,
                                               lambda _: 5, ],
-                                             a)
+                                                    a)
             return jax.vmap(f)(operands)
 
-        r = f(bst.random.randint(-20, 20, 200))
+        r = f(brainstate.random.randint(-20, 20, 200))
         self.assertTrue(r.size == 200)
 
     def test_grad1(self):
         def F2(x):
-            return bst.compile.ifelse((x >= 10, x < 10),
-                                      [lambda x: x, lambda x: x ** 2, ],
-                                      x)
+            return brainstate.compile.ifelse((x >= 10, x < 10),
+                                             [lambda x: x, lambda x: x ** 2, ],
+                                             x)
 
         self.assertTrue(jax.grad(F2)(9.0) == 18.)
         self.assertTrue(jax.grad(F2)(11.0) == 1.)
 
     def test_grad2(self):
         def F3(x):
-            return bst.compile.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
-                                      [lambda x: x,
+            return brainstate.compile.ifelse((x >= 10, jnp.logical_and(x >= 0, x < 10), x < 0),
+                                             [lambda x: x,
                                        lambda x: x ** 2,
                                        lambda x: x ** 4, ],
-                                      x)
+                                             x)
 
         self.assertTrue(jax.grad(F3)(9.0) == 18.)
         self.assertTrue(jax.grad(F3)(11.0) == 1.)

brainstate/compile/_error_if_test.py

@@ -13,43 +13,40 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax
 import jax.numpy as jnp
-import jaxlib.xla_extension
 
-import brainstate as bst
+import brainstate
 
 
 class TestJitError(unittest.TestCase):
     def test1(self):
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-            bst.compile.jit_error_if(True, 'error')
+        with self.assertRaises(Exception):
+            brainstate.compile.jit_error_if(True, 'error')
 
         def err_f(x):
             raise ValueError(f'error: {x}')
 
-        bst.compile.jit_error_if(False, err_f, 1.)
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
-            bst.compile.jit_error_if(True, err_f, 1.)
+        brainstate.compile.jit_error_if(False, err_f, 1.)
+        with self.assertRaises(Exception):
+            brainstate.compile.jit_error_if(True, err_f, 1.)
 
     def test_vmap(self):
         def f(x):
-            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+            brainstate.compile.jit_error_if(x, 'error: {x}', x=x)
 
         jax.vmap(f)(jnp.array([False, False, False]))
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        with self.assertRaises(Exception):
             jax.vmap(f)(jnp.array([True, False, False]))
 
     def test_vmap_vmap(self):
         def f(x):
-            bst.compile.jit_error_if(x, 'error: {x}', x=x)
+            brainstate.compile.jit_error_if(x, 'error: {x}', x=x)
 
         jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
                                          [False, False, False]]))
-        with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+        with self.assertRaises(Exception):
             jax.vmap(jax.vmap(f))(jnp.array([[False, False, False],
                                              [True, False, False]]))