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

brainstate/compile/_loop_collect_return_test.py

@@ -13,20 +13,18 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax.numpy as jnp
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestForLoop(unittest.TestCase):
     def test_for_loop(self):
-        a = bst.ShortTermState(0.)
-        b = bst.ShortTermState(0.)
+        a = brainstate.ShortTermState(0.)
+        b = brainstate.ShortTermState(0.)
 
         def f(i):
             a.value += (1 + b.value)
@@ -34,7 +32,7 @@ class TestForLoop(unittest.TestCase):
 
         n_iter = 10
         ops = np.arange(n_iter)
-        r = bst.compile.for_loop(f, ops)
+        r = brainstate.compile.for_loop(f, ops)
 
         print(a)
         print(b)
@@ -42,8 +40,8 @@ class TestForLoop(unittest.TestCase):
         self.assertTrue(jnp.allclose(r, ops + 1))
 
     def test_checkpointed_for_loop(self):
-        a = bst.ShortTermState(0.)
-        b = bst.ShortTermState(0.)
+        a = brainstate.ShortTermState(0.)
+        b = brainstate.ShortTermState(0.)
 
         def f(i):
             a.value += (1 + b.value)
@@ -51,7 +49,7 @@ class TestForLoop(unittest.TestCase):
 
         n_iter = 18
         ops = jnp.arange(n_iter)
-        r = bst.compile.checkpointed_for_loop(f, ops, base=2, pbar=bst.compile.ProgressBar())
+        r = brainstate.compile.checkpointed_for_loop(f, ops, base=2, pbar=brainstate.compile.ProgressBar())
 
         print(a)
         print(b)

brainstate/compile/_loop_no_collection_test.py

@@ -13,17 +13,16 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 from unittest import TestCase
 
-import brainstate as bst
+import brainstate 
 
 
 class TestWhileLoop(TestCase):
     def test1(self):
-        a = bst.State(1.)
-        b = bst.State(20.)
+        a = brainstate.State(1.)
+        b = brainstate.State(20.)
 
         def cond(_):
             return a.value < b.value
@@ -31,13 +30,13 @@ class TestWhileLoop(TestCase):
         def body(_):
             a.value += 1.
 
-        bst.compile.while_loop(cond, body, None)
+        brainstate.compile.while_loop(cond, body, None)
 
         print(a.value, b.value)
 
     def test2(self):
-        a = bst.State(1.)
-        b = bst.State(20.)
+        a = brainstate.State(1.)
+        b = brainstate.State(20.)
 
         def cond(x):
             return a.value < b.value
@@ -46,6 +45,6 @@ class TestWhileLoop(TestCase):
             a.value += x
             return x
 
-        r = bst.compile.while_loop(cond, body, 1.)
+        r = brainstate.compile.while_loop(cond, body, 1.)
 
         print(a.value, b.value, r)

brainstate/compile/_make_jaxpr.py

@@ -62,8 +62,8 @@ import jax
 from jax._src import source_info_util
 from jax._src.linear_util import annotate
 from jax._src.traceback_util import api_boundary
-from jax.api_util import shaped_abstractify, debug_info
-from jax.extend.linear_util import transformation_with_aux, wrap_init
+from jax.api_util import shaped_abstractify
+from jax.extend.linear_util import transformation_with_aux
 from jax.interpreters import partial_eval as pe
 
 from brainstate._compatible_import import (
@@ -73,6 +73,7 @@ from brainstate._compatible_import import (
     safe_zip,
     unzip2,
     wraps,
+    wrap_init,
 )
 from brainstate._state import State, StateTraceStack
 from brainstate._utils import set_module_as
@@ -96,7 +97,7 @@ def _ensure_index_tuple(x: Any) -> tuple[int, ...]:
         return tuple(safe_map(operator.index, x))
 
 
-def _new_arg_fn(frame, trace, aval):
+def _jax_v04_new_arg_fn(frame, trace, aval):
     """
     Transform a new argument to a tracer.
 
@@ -117,27 +118,41 @@ def _new_arg_fn(frame, trace, aval):
     return tracer
 
 
-def _new_jax_trace():
+def _jax_v04_new_jax_trace():
     main = jax.core.thread_local_state.trace_state.trace_stack.stack[-1]
     frame = main.jaxpr_stack[-1]
     trace = pe.DynamicJaxprTrace(main, jax.core.cur_sublevel())
     return frame, trace
 
 
+def _jax_v04_new_arg():
+    # Should be within the calling of ``jax.make_jaxpr()``
+    frame, trace = _jax_v04_new_jax_trace()
+    # Set the function to transform the new argument to a tracer
+    fn = functools.partial(_jax_v04_new_arg_fn, frame, trace)
+    return fn
+
+
+def _jax_new_version_new_arg():
+    trace = jax.core.trace_ctx.trace
+
+    def wrapper(x):
+        if jax.__version_info__ < (0, 6, 1):
+            return trace.new_arg(shaped_abstractify(x))
+        else:
+            return trace.new_arg(shaped_abstractify(x), source_info=source_info_util.current())
+
+    return wrapper
+
+
 def _init_state_trace_stack(name) -> StateTraceStack:
     state_trace: StateTraceStack = StateTraceStack(name=name)
 
     if jax.__version_info__ < (0, 4, 36):
-        # Should be within the calling of ``jax.make_jaxpr()``
-        frame, trace = _new_jax_trace()
-        # Set the function to transform the new argument to a tracer
-        state_trace.set_new_arg(functools.partial(_new_arg_fn, frame, trace))
-        return state_trace
-
+        state_trace.set_new_arg(_jax_v04_new_arg())
     else:
-        trace = jax.core.trace_ctx.trace
-        state_trace.set_new_arg(trace.new_arg)
-        return state_trace
+        state_trace.set_new_arg(_jax_new_version_new_arg())
+    return state_trace
 
 
 class StatefulFunction(PrettyObject):
@@ -743,7 +758,7 @@ def _make_jaxpr(
     @wraps(fun)
     @api_boundary
     def make_jaxpr_f(*args, **kwargs):
-        f = wrap_init(fun, debug_info=debug_info('make_jaxpr', fun, args, kwargs))
+        f = wrap_init(fun, (), {}, 'brainstate.compile.make_jaxpr')
         if static_argnums:
             dyn_argnums = [i for i in range(len(args)) if i not in static_argnums]
             f, args = jax.api_util.argnums_partial(f, dyn_argnums, args)

brainstate/compile/_make_jaxpr_test.py

@@ -21,7 +21,7 @@ import jax
 import jax.numpy as jnp
 import pytest
 
-import brainstate as bst
+import brainstate
 from brainstate._compatible_import import jaxpr_as_fun
 
 
@@ -29,10 +29,10 @@ class TestMakeJaxpr(unittest.TestCase):
     def test_compar_jax_make_jaxpr(self):
         def func4(arg):  # Arg is a pair
             temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = bst.random.rand_like(arg[0])
+            c = brainstate.random.rand_like(arg[0])
             return jnp.sum(temp + c)
 
-        key = bst.random.DEFAULT.value
+        key = brainstate.random.DEFAULT.value
         jaxpr = jax.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
         print(jaxpr)
         self.assertTrue(len(jaxpr.in_avals) == 2)
@@ -40,66 +40,66 @@ class TestMakeJaxpr(unittest.TestCase):
         self.assertTrue(len(jaxpr.out_avals) == 1)
         self.assertTrue(jnp.allclose(jaxpr.consts[0], key))
 
-        bst.random.seed(1)
-        print(bst.random.DEFAULT.value)
+        brainstate.random.seed(1)
+        print(brainstate.random.DEFAULT.value)
 
-        jaxpr2, states = bst.compile.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
+        jaxpr2, states = brainstate.compile.make_jaxpr(func4)((jnp.zeros(8), jnp.ones(8)))
         print(jaxpr2)
         self.assertTrue(len(jaxpr2.in_avals) == 3)
         self.assertTrue(len(jaxpr2.out_avals) == 2)
         self.assertTrue(len(jaxpr2.consts) == 0)
-        print(bst.random.DEFAULT.value)
+        print(brainstate.random.DEFAULT.value)
 
     def test_StatefulFunction_1(self):
         def func4(arg):  # Arg is a pair
             temp = arg[0] + jnp.sin(arg[1]) * 3.
-            c = bst.random.rand_like(arg[0])
+            c = brainstate.random.rand_like(arg[0])
             return jnp.sum(temp + c)
 
-        fun = bst.compile.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
+        fun = brainstate.compile.StatefulFunction(func4).make_jaxpr((jnp.zeros(8), jnp.ones(8)))
         print(fun.get_states())
         print(fun.get_jaxpr())
 
     def test_StatefulFunction_2(self):
-        st1 = bst.State(jnp.ones(10))
+        st1 = brainstate.State(jnp.ones(10))
 
         def f1(x):
             st1.value = x + st1.value
 
         def f2(x):
-            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
             c = 1. + x
             return c
 
         def f3(x):
-            jaxpr = bst.compile.make_jaxpr(f1)(x)
+            jaxpr = brainstate.compile.make_jaxpr(f1)(x)
             c = 1.
             return c
 
         print()
-        jaxpr = bst.compile.make_jaxpr(f1)(jnp.zeros(1))
+        jaxpr = brainstate.compile.make_jaxpr(f1)(jnp.zeros(1))
         print(jaxpr)
         jaxpr = jax.make_jaxpr(f2)(jnp.zeros(1))
         print(jaxpr)
         jaxpr = jax.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
-        jaxpr, _ = bst.compile.make_jaxpr(f3)(jnp.zeros(1))
+        jaxpr, _ = brainstate.compile.make_jaxpr(f3)(jnp.zeros(1))
         print(jaxpr)
         self.assertTrue(jnp.allclose(jaxpr_as_fun(jaxpr)(jnp.zeros(1), st1.value)[0],
                                      f3(jnp.zeros(1))))
 
     def test_compar_jax_make_jaxpr2(self):
-        st1 = bst.State(jnp.ones(10))
+        st1 = brainstate.State(jnp.ones(10))
 
         def fa(x):
             st1.value = x + st1.value
 
         def ffa(x):
-            jaxpr, states = bst.compile.make_jaxpr(fa)(x)
+            jaxpr, states = brainstate.compile.make_jaxpr(fa)(x)
             c = 1. + x
             return c
 
-        jaxpr, states = bst.compile.make_jaxpr(ffa)(jnp.zeros(1))
+        jaxpr, states = brainstate.compile.make_jaxpr(ffa)(jnp.zeros(1))
         print()
         print(jaxpr)
         print(states)
@@ -112,7 +112,7 @@ class TestMakeJaxpr(unittest.TestCase):
         def fa(x):
             return 1.
 
-        jaxpr, states = bst.compile.make_jaxpr(fa)(jnp.zeros(1))
+        jaxpr, states = brainstate.compile.make_jaxpr(fa)(jnp.zeros(1))
         print()
         print(jaxpr)
         print(states)
@@ -125,9 +125,9 @@ class TestMakeJaxpr(unittest.TestCase):
 def test_return_states():
     import jax.numpy
 
-    a = bst.State(jax.numpy.ones(3))
+    a = brainstate.State(jax.numpy.ones(3))
 
-    @bst.compile.jit
+    @brainstate.compile.jit
     def f():
         return a
 

brainstate/functional/_activations_test.py

@@ -25,48 +25,48 @@ from absl.testing import parameterized
 from jax._src import test_util as jtu
 from jax.test_util import check_grads
 
-import brainstate as bst
+import brainstate
 
 
 class NNFunctionsTest(jtu.JaxTestCase):
     @jtu.skip_on_flag("jax_skip_slow_tests", True)
     def testSoftplusGrad(self):
-        check_grads(bst.functional.softplus, (1e-8,), order=4, )
+        check_grads(brainstate.functional.softplus, (1e-8,), order=4, )
 
     def testSoftplusGradZero(self):
-        check_grads(bst.functional.softplus, (0.,), order=1)
+        check_grads(brainstate.functional.softplus, (0.,), order=1)
 
     def testSoftplusGradInf(self):
-        self.assertAllClose(1., jax.grad(bst.functional.softplus)(float('inf')))
+        self.assertAllClose(1., jax.grad(brainstate.functional.softplus)(float('inf')))
 
     def testSoftplusGradNegInf(self):
-        check_grads(bst.functional.softplus, (-float('inf'),), order=1)
+        check_grads(brainstate.functional.softplus, (-float('inf'),), order=1)
 
     def testSoftplusGradNan(self):
-        check_grads(bst.functional.softplus, (float('nan'),), order=1)
+        check_grads(brainstate.functional.softplus, (float('nan'),), order=1)
 
     @parameterized.parameters([int, float] + jtu.dtypes.floating + jtu.dtypes.integer)
     def testSoftplusZero(self, dtype):
-        self.assertEqual(jnp.log(dtype(2)), bst.functional.softplus(dtype(0)))
+        self.assertEqual(jnp.log(dtype(2)), brainstate.functional.softplus(dtype(0)))
 
     def testSparseplusGradZero(self):
-        check_grads(bst.functional.sparse_plus, (-2.,), order=1)
+        check_grads(brainstate.functional.sparse_plus, (-2.,), order=1)
 
     def testSparseplusGrad(self):
-        check_grads(bst.functional.sparse_plus, (0.,), order=1)
+        check_grads(brainstate.functional.sparse_plus, (0.,), order=1)
 
     def testSparseplusAndSparseSigmoid(self):
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(0.),
-            bst.functional.sparse_sigmoid(0.),
+            jax.grad(brainstate.functional.sparse_plus)(0.),
+            brainstate.functional.sparse_sigmoid(0.),
             check_dtypes=False)
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(2.),
-            bst.functional.sparse_sigmoid(2.),
+            jax.grad(brainstate.functional.sparse_plus)(2.),
+            brainstate.functional.sparse_sigmoid(2.),
             check_dtypes=False)
         self.assertAllClose(
-            jax.grad(bst.functional.sparse_plus)(-2.),
-            bst.functional.sparse_sigmoid(-2.),
+            jax.grad(brainstate.functional.sparse_plus)(-2.),
+            brainstate.functional.sparse_sigmoid(-2.),
             check_dtypes=False)
 
     #   def testSquareplusGrad(self):
@@ -107,55 +107,55 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     @parameterized.parameters([float] + jtu.dtypes.floating)
     def testMishZero(self, dtype):
-        self.assertEqual(dtype(0), bst.functional.mish(dtype(0)))
+        self.assertEqual(dtype(0), brainstate.functional.mish(dtype(0)))
 
     def testReluGrad(self):
         rtol = None
-        check_grads(bst.functional.relu, (1.,), order=3, rtol=rtol)
-        check_grads(bst.functional.relu, (-1.,), order=3, rtol=rtol)
-        jaxpr = jax.make_jaxpr(jax.grad(bst.functional.relu))(0.)
+        check_grads(brainstate.functional.relu, (1.,), order=3, rtol=rtol)
+        check_grads(brainstate.functional.relu, (-1.,), order=3, rtol=rtol)
+        jaxpr = jax.make_jaxpr(jax.grad(brainstate.functional.relu))(0.)
         self.assertGreaterEqual(len(jaxpr.jaxpr.eqns), 2)
 
     def testRelu6Grad(self):
         rtol = None
-        check_grads(bst.functional.relu6, (1.,), order=3, rtol=rtol)
-        check_grads(bst.functional.relu6, (-1.,), order=3, rtol=rtol)
-        self.assertAllClose(jax.grad(bst.functional.relu6)(0.), 0., check_dtypes=False)
-        self.assertAllClose(jax.grad(bst.functional.relu6)(6.), 0., check_dtypes=False)
+        check_grads(brainstate.functional.relu6, (1.,), order=3, rtol=rtol)
+        check_grads(brainstate.functional.relu6, (-1.,), order=3, rtol=rtol)
+        self.assertAllClose(jax.grad(brainstate.functional.relu6)(0.), 0., check_dtypes=False)
+        self.assertAllClose(jax.grad(brainstate.functional.relu6)(6.), 0., check_dtypes=False)
 
     def testSoftplusValue(self):
-        val = bst.functional.softplus(89.)
+        val = brainstate.functional.softplus(89.)
         self.assertAllClose(val, 89., check_dtypes=False)
 
     def testSparseplusValue(self):
-        val = bst.functional.sparse_plus(89.)
+        val = brainstate.functional.sparse_plus(89.)
         self.assertAllClose(val, 89., check_dtypes=False)
 
     def testSparsesigmoidValue(self):
-        self.assertAllClose(bst.functional.sparse_sigmoid(-2.), 0., check_dtypes=False)
-        self.assertAllClose(bst.functional.sparse_sigmoid(2.), 1., check_dtypes=False)
-        self.assertAllClose(bst.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(-2.), 0., check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(2.), 1., check_dtypes=False)
+        self.assertAllClose(brainstate.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
 
     #   def testSquareplusValue(self):
     #     val = bst.functional.squareplus(1e3)
     #     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testMishValue(self):
-        val = bst.functional.mish(1e3)
+        val = brainstate.functional.mish(1e3)
         self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testEluValue(self):
-        val = bst.functional.elu(1e4)
+        val = brainstate.functional.elu(1e4)
         self.assertAllClose(val, 1e4, check_dtypes=False)
 
     def testGluValue(self):
-        val = bst.functional.glu(jnp.array([1.0, 0.0]), axis=0)
+        val = brainstate.functional.glu(jnp.array([1.0, 0.0]), axis=0)
         self.assertAllClose(val, jnp.array([0.5]))
 
     @parameterized.parameters(False, True)
     def testGeluIntType(self, approximate):
-        val_float = bst.functional.gelu(jnp.array(-1.0), approximate=approximate)
-        val_int = bst.functional.gelu(jnp.array(-1), approximate=approximate)
+        val_float = brainstate.functional.gelu(jnp.array(-1.0), approximate=approximate)
+        val_int = brainstate.functional.gelu(jnp.array(-1), approximate=approximate)
         self.assertAllClose(val_float, val_int)
 
     @parameterized.parameters(False, True)
@@ -166,19 +166,19 @@ class NNFunctionsTest(jtu.JaxTestCase):
         rng = jtu.rand_default(self.rng())
         args_maker = lambda: [rng((4, 5, 6), jnp.float32)]
         self._CheckAgainstNumpy(
-            gelu_reference, partial(bst.functional.gelu, approximate=approximate), args_maker,
+            gelu_reference, partial(brainstate.functional.gelu, approximate=approximate), args_maker,
             check_dtypes=False, tol=1e-3 if approximate else None)
 
     @parameterized.parameters(*itertools.product(
         (jnp.float32, jnp.bfloat16, jnp.float16),
-        (partial(bst.functional.gelu, approximate=False),
-         partial(bst.functional.gelu, approximate=True),
-         bst.functional.relu,
-         bst.functional.softplus,
-         bst.functional.sparse_plus,
-         bst.functional.sigmoid,
+        (partial(brainstate.functional.gelu, approximate=False),
+         partial(brainstate.functional.gelu, approximate=True),
+         brainstate.functional.relu,
+         brainstate.functional.softplus,
+         brainstate.functional.sparse_plus,
+         brainstate.functional.sigmoid,
          #  bst.functional.squareplus,
-         bst.functional.mish)))
+         brainstate.functional.mish)))
     def testDtypeMatchesInput(self, dtype, fn):
         x = jnp.zeros((), dtype=dtype)
         out = fn(x)
@@ -187,26 +187,26 @@ class NNFunctionsTest(jtu.JaxTestCase):
     def testEluMemory(self):
         # see https://github.com/google/jax/pull/1640
         with jax.enable_checks(False):  # With checks we materialize the array
-            jax.make_jaxpr(lambda: bst.functional.elu(jnp.ones((10 ** 12,))))  # don't oom
+            jax.make_jaxpr(lambda: brainstate.functional.elu(jnp.ones((10 ** 12,))))  # don't oom
 
     def testHardTanhMemory(self):
         # see https://github.com/google/jax/pull/1640
         with jax.enable_checks(False):  # With checks we materialize the array
-            jax.make_jaxpr(lambda: bst.functional.hard_tanh(jnp.ones((10 ** 12,))))  # don't oom
+            jax.make_jaxpr(lambda: brainstate.functional.hard_tanh(jnp.ones((10 ** 12,))))  # don't oom
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxEmptyArray(self, fn):
         x = jnp.array([], dtype=float)
         self.assertArraysEqual(fn(x), x)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxEmptyMask(self, fn):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.zeros_like(x, dtype=bool)
-        expected = jnp.full_like(x, 0.0 if fn is bst.functional.softmax else -jnp.inf)
+        expected = jnp.full_like(x, 0.0 if fn is brainstate.functional.softmax else -jnp.inf)
         self.assertArraysEqual(fn(x, where=m), expected)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxWhereMask(self, fn):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.array([True, False, True, True])
@@ -214,10 +214,10 @@ class NNFunctionsTest(jtu.JaxTestCase):
         out = fn(x, where=m)
         self.assertAllClose(out[m], fn(x[m]))
 
-        probs = out if fn is bst.functional.softmax else jnp.exp(out)
+        probs = out if fn is brainstate.functional.softmax else jnp.exp(out)
         self.assertAllClose(probs.sum(), 1.0)
 
-    @parameterized.parameters([bst.functional.softmax, bst.functional.log_softmax])
+    @parameterized.parameters([brainstate.functional.softmax, brainstate.functional.log_softmax])
     def testSoftmaxWhereGrad(self, fn):
         # regression test for https://github.com/google/jax/issues/19490
         x = jnp.array([36., 10000.])
@@ -229,46 +229,46 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     def testSoftmaxGrad(self):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
-        jtu.check_grads(bst.functional.softmax, (x,), order=2, atol=5e-3)
+        jtu.check_grads(brainstate.functional.softmax, (x,), order=2, atol=5e-3)
 
     def testStandardizeWhereMask(self):
         x = jnp.array([5.5, 1.3, -4.2, 0.9])
         m = jnp.array([True, False, True, True])
         x_filtered = jnp.take(x, jnp.array([0, 2, 3]))
 
-        out_masked = jnp.take(bst.functional.standardize(x, where=m), jnp.array([0, 2, 3]))
-        out_filtered = bst.functional.standardize(x_filtered)
+        out_masked = jnp.take(brainstate.functional.standardize(x, where=m), jnp.array([0, 2, 3]))
+        out_filtered = brainstate.functional.standardize(x_filtered)
 
         self.assertAllClose(out_masked, out_filtered)
 
     def testOneHot(self):
-        actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([0, 1, 2]), 3)
         expected = jnp.array([[1., 0., 0.],
                               [0., 1., 0.],
                               [0., 0., 1.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3)
         expected = jnp.array([[0., 1., 0.],
                               [0., 0., 1.],
                               [1., 0., 0.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotOutOfBound(self):
-        actual = bst.functional.one_hot(jnp.array([-1, 3]), 3)
+        actual = brainstate.functional.one_hot(jnp.array([-1, 3]), 3)
         expected = jnp.array([[0., 0., 0.],
                               [0., 0., 0.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotNonArrayInput(self):
-        actual = bst.functional.one_hot([0, 1, 2], 3)
+        actual = brainstate.functional.one_hot([0, 1, 2], 3)
         expected = jnp.array([[1., 0., 0.],
                               [0., 1., 0.],
                               [0., 0., 1.]])
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testOneHotCustomDtype(self):
-        actual = bst.functional.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
+        actual = brainstate.functional.one_hot(jnp.array([0, 1, 2]), 3, dtype=jnp.bool_)
         expected = jnp.array([[True, False, False],
                               [False, True, False],
                               [False, False, True]])
@@ -279,14 +279,14 @@ class NNFunctionsTest(jtu.JaxTestCase):
                               [0., 0., 1.],
                               [1., 0., 0.]]).T
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=0)
         self.assertAllClose(actual, expected, check_dtypes=False)
 
-        actual = bst.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
+        actual = brainstate.functional.one_hot(jnp.array([1, 2, 0]), 3, axis=-2)
         self.assertAllClose(actual, expected, check_dtypes=False)
 
     def testTanhExists(self):
-        print(bst.functional.tanh)  # doesn't crash
+        print(brainstate.functional.tanh)  # doesn't crash
 
     def testCustomJVPLeak(self):
         # https://github.com/google/jax/issues/8171
@@ -295,7 +295,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
             a = jnp.array(1.)
 
             def f(hx, _):
-                hx = bst.functional.sigmoid(hx + a)
+                hx = brainstate.functional.sigmoid(hx + a)
                 return hx, None
 
             hx = jnp.array(0.)