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

brainstate/nn/{_interaction/_linear_test.py → _linear_test.py}

@@ -14,14 +14,12 @@
 # ==============================================================================
 
 
-from __future__ import annotations
-
 import unittest
 
 import brainunit as u
 from absl.testing import parameterized
 
-import brainstate as bst
+import brainstate
 
 
 class TestDense(parameterized.TestCase):
@@ -32,19 +30,19 @@ class TestDense(parameterized.TestCase):
         num_out=[20, ]
     )
     def test_Dense1(self, size, num_out):
-        f = bst.nn.Linear(10, num_out)
-        x = bst.random.random(size)
+        f = brainstate.nn.Linear(10, num_out)
+        x = brainstate.random.random(size)
         y = f(x)
         self.assertTrue(y.shape == size[:-1] + (num_out,))
 
 
 class TestSparseMatrix(unittest.TestCase):
     def test_csr(self):
-        data = bst.random.rand(10, 20)
+        data = brainstate.random.rand(10, 20)
         data = data * (data > 0.9)
-        f = bst.nn.SparseLinear(u.sparse.CSR.fromdense(data))
+        f = brainstate.nn.SparseLinear(u.sparse.CSR.fromdense(data))
 
-        x = bst.random.rand(10)
+        x = brainstate.random.rand(10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(
@@ -53,7 +51,7 @@ class TestSparseMatrix(unittest.TestCase):
             )
         )
 
-        x = bst.random.rand(5, 10)
+        x = brainstate.random.rand(5, 10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(
@@ -63,11 +61,11 @@ class TestSparseMatrix(unittest.TestCase):
         )
 
     def test_csc(self):
-        data = bst.random.rand(10, 20)
+        data = brainstate.random.rand(10, 20)
         data = data * (data > 0.9)
-        f = bst.nn.SparseLinear(u.sparse.CSC.fromdense(data))
+        f = brainstate.nn.SparseLinear(u.sparse.CSC.fromdense(data))
 
-        x = bst.random.rand(10)
+        x = brainstate.random.rand(10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(
@@ -76,7 +74,7 @@ class TestSparseMatrix(unittest.TestCase):
             )
         )
 
-        x = bst.random.rand(5, 10)
+        x = brainstate.random.rand(5, 10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(
@@ -86,11 +84,11 @@ class TestSparseMatrix(unittest.TestCase):
         )
 
     def test_coo(self):
-        data = bst.random.rand(10, 20)
+        data = brainstate.random.rand(10, 20)
         data = data * (data > 0.9)
-        f = bst.nn.SparseLinear(u.sparse.COO.fromdense(data))
+        f = brainstate.nn.SparseLinear(u.sparse.COO.fromdense(data))
 
-        x = bst.random.rand(10)
+        x = brainstate.random.rand(10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(
@@ -99,7 +97,7 @@ class TestSparseMatrix(unittest.TestCase):
             )
         )
 
-        x = bst.random.rand(5, 10)
+        x = brainstate.random.rand(5, 10)
         y = f(x)
         self.assertTrue(
             u.math.allclose(

brainstate/nn/{_event/__init__.py → _ltp.py}

@@ -16,11 +16,13 @@
 # -*- coding: utf-8 -*-
 
 
-from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
-from ._linear_mv import EventLinear
+from ._synapse import Synapse
 
 __all__ = [
-    'EventLinear',
-    'EventFixedProb',
-    'EventFixedNumConn',
+    'LongTermPlasticity',
 ]
+
+
+class LongTermPlasticity(Synapse):
+    pass
+

brainstate/nn/_module_test.py

@@ -13,20 +13,17 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 import unittest
 
 import jax.numpy as jnp
-import jaxlib.xla_extension
 
-import brainstate as bst
+import brainstate
 
 
 class TestDelay(unittest.TestCase):
     def test_delay1(self):
-        a = bst.State(bst.random.random(10, 20))
-        delay = bst.nn.Delay(a.value)
+        a = brainstate.State(brainstate.random.random(10, 20))
+        delay = brainstate.nn.Delay(a.value)
         delay.register_entry('a', 1.)
         delay.register_entry('b', 2.)
         delay.register_entry('c', None)
@@ -36,7 +33,7 @@ class TestDelay(unittest.TestCase):
             delay.register_entry('c', 10.)
 
     def test_rotation_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones((1,)))
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -53,7 +50,7 @@ class TestDelay(unittest.TestCase):
         # print(rotation_delay.max_length)
 
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             rotation_delay.update(jnp.ones((1,)) * i)
             # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
             self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
@@ -61,7 +58,7 @@ class TestDelay(unittest.TestCase):
             self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
 
     def test_concat_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones([1]), delay_method='concat')
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -74,7 +71,7 @@ class TestDelay(unittest.TestCase):
 
         print()
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             rotation_delay.update(jnp.ones((1,)) * i)
             print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
             self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
@@ -83,40 +80,40 @@ class TestDelay(unittest.TestCase):
         # bst.util.clear_buffer_memory()
 
     def test_jit_erro(self):
-        rotation_delay = bst.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
         rotation_delay.init_state()
 
-        with bst.environ.context(i=0, t=0, jit_error_check=True):
+        with brainstate.environ.context(i=0, t=0, jit_error_check=True):
             rotation_delay.retrieve_at_time(-2.0)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(-2.1)
             rotation_delay.retrieve_at_time(-2.01)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(-2.09)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(0.1)
-            with self.assertRaises(jaxlib.xla_extension.XlaRuntimeError):
+            with self.assertRaises(Exception):
                 rotation_delay.retrieve_at_time(0.01)
 
     def test_round_interp(self):
         for shape in [(1,), (1, 1), (1, 1, 1)]:
             for delay_method in ['rotation', 'concat']:
-                rotation_delay = bst.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                              interp_method='round')
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='round')
                 t0, n1 = 0.01, 0
                 t1, n1 = 1.04, 10
                 t2, n2 = 1.06, 11
                 rotation_delay.init_state()
 
-                @bst.compile.jit
+                @brainstate.compile.jit
                 def retrieve(td, i):
-                    with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
                         return rotation_delay.retrieve_at_time(td)
 
                 print()
                 for i in range(100):
-                    t = i * bst.environ.get_dt()
-                    with bst.environ.context(i=i, t=t):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
                         rotation_delay.update(jnp.ones(shape) * i)
                         print(i,
                               retrieve(t - t0, i),
@@ -131,22 +128,22 @@ class TestDelay(unittest.TestCase):
             for delay_method in ['rotation', 'concat']:
                 print(shape, delay_method)
 
-                rotation_delay = bst.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                              interp_method='linear_interp')
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='linear_interp')
                 t0, n0 = 0.01, 0.1
                 t1, n1 = 1.04, 10.4
                 t2, n2 = 1.06, 10.6
                 rotation_delay.init_state()
 
-                @bst.compile.jit
+                @brainstate.compile.jit
                 def retrieve(td, i):
-                    with bst.environ.context(i=i, t=i * bst.environ.get_dt()):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
                         return rotation_delay.retrieve_at_time(td)
 
                 print()
                 for i in range(100):
-                    t = i * bst.environ.get_dt()
-                    with bst.environ.context(i=i, t=t):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
                         rotation_delay.update(jnp.ones(shape) * i)
                         print(i,
                               retrieve(t - t0, i),
@@ -157,8 +154,8 @@ class TestDelay(unittest.TestCase):
                         self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
 
     def test_rotation_and_concat_delay(self):
-        rotation_delay = bst.nn.Delay(jnp.ones((1,)))
-        concat_delay = bst.nn.Delay(jnp.ones([1]), delay_method='concat')
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
+        concat_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
         t0 = 0.
         t1, n1 = 1., 10
         t2, n2 = 2., 20
@@ -175,7 +172,7 @@ class TestDelay(unittest.TestCase):
 
         print()
         for i in range(100):
-            bst.environ.set(i=i)
+            brainstate.environ.set(i=i)
             new = jnp.ones((1,)) * i
             rotation_delay.update(new)
             concat_delay.update(new)
@@ -186,17 +183,17 @@ class TestDelay(unittest.TestCase):
 
 class TestModule(unittest.TestCase):
     def test_states(self):
-        class A(bst.nn.Module):
+        class A(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
-                self.a = bst.State(bst.random.random(10, 20))
-                self.b = bst.State(bst.random.random(10, 20))
+                self.a = brainstate.State(brainstate.random.random(10, 20))
+                self.b = brainstate.State(brainstate.random.random(10, 20))
 
-        class B(bst.nn.Module):
+        class B(brainstate.nn.Module):
             def __init__(self):
                 super().__init__()
                 self.a = A()
-                self.b = bst.State(bst.random.random(10, 20))
+                self.b = brainstate.State(brainstate.random.random(10, 20))
 
         b = B()
         print()
@@ -207,5 +204,5 @@ class TestModule(unittest.TestCase):
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/{_dyn_impl/_dynamics_neuron.py → _neuron.py}

@@ -22,9 +22,9 @@ import jax
 
 from brainstate import init, surrogate, environ
 from brainstate._state import HiddenState, ShortTermState
-from brainstate.nn._dynamics._dynamics_base import Dynamics
-from brainstate.nn._exp_euler import exp_euler_step
 from brainstate.typing import ArrayLike, Size
+from ._dynamics import Dynamics
+from ._exp_euler import exp_euler_step
 
 __all__ = [
     'Neuron', 'IF', 'LIF', 'LIFRef', 'ALIF',

brainstate/nn/{_dyn_impl/_dynamics_neuron_test.py → _neuron_test.py}

@@ -15,7 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
 
 import unittest
 
@@ -23,7 +22,7 @@ import brainunit as u
 import jax
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 from brainstate.nn import IF, LIF, ALIF
 
 
@@ -35,13 +34,13 @@ class TestNeuron(unittest.TestCase):
 
     def test_neuron_base_class(self):
         with self.assertRaises(NotImplementedError):
-            bst.nn.Neuron(self.in_size).get_spike()  # Neuron is an abstract base class
+            brainstate.nn.Neuron(self.in_size).get_spike()  # Neuron is an abstract base class
 
     def generate_input(self):
-        return bst.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mA
+        return brainstate.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mA
 
     def test_if_neuron(self):
-        with bst.environ.context(dt=0.1 * u.ms):
+        with brainstate.environ.context(dt=0.1 * u.ms):
             neuron = IF(self.in_size)
             inputs = self.generate_input()
 
@@ -62,7 +61,7 @@ class TestNeuron(unittest.TestCase):
             self.assertTrue(jnp.all((spikes >= 0) & (spikes <= 1)))
 
     def test_lif_neuron(self):
-        with bst.environ.context(dt=0.1 * u.ms):
+        with brainstate.environ.context(dt=0.1 * u.ms):
             tau = 20.0 * u.ms
             neuron = LIF(self.in_size, tau=tau)
             inputs = self.generate_input()
@@ -74,7 +73,7 @@ class TestNeuron(unittest.TestCase):
 
             # Test forward pass
             state = neuron.init_state(self.batch_size)
-            call = bst.compile.jit(neuron)
+            call = brainstate.compile.jit(neuron)
 
             for t in range(self.time_steps):
                 out = call(inputs[t])
@@ -94,8 +93,8 @@ class TestNeuron(unittest.TestCase):
 
         # Test forward pass
         neuron.init_state(self.batch_size)
-        call = bst.compile.jit(neuron)
-        with bst.environ.context(dt=0.1 * u.ms):
+        call = brainstate.compile.jit(neuron)
+        with brainstate.environ.context(dt=0.1 * u.ms):
             for t in range(self.time_steps):
                 out = call(inputs[t])
                 self.assertEqual(out.shape, (self.batch_size, self.in_size))
@@ -113,8 +112,8 @@ class TestNeuron(unittest.TestCase):
             neuron = NeuronClass(self.in_size, spk_reset='soft')
             inputs = self.generate_input()
             state = neuron.init_state(self.batch_size)
-            call = bst.compile.jit(neuron)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(neuron)
+            with brainstate.environ.context(dt=0.1 * u.ms):
                 for t in range(self.time_steps):
                     out = call(inputs[t])
                     self.assertTrue(jnp.all(neuron.V.value <= neuron.V_th))
@@ -124,8 +123,8 @@ class TestNeuron(unittest.TestCase):
             neuron = NeuronClass(self.in_size, spk_reset='hard')
             inputs = self.generate_input()
             state = neuron.init_state(self.batch_size)
-            call = bst.compile.jit(neuron)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(neuron)
+            with brainstate.environ.context(dt=0.1 * u.ms):
                 for t in range(self.time_steps):
                     out = call(inputs[t])
                     self.assertTrue(jnp.all((neuron.V.value < neuron.V_th) | (neuron.V.value == 0. * u.mV)))
@@ -135,8 +134,8 @@ class TestNeuron(unittest.TestCase):
             neuron = NeuronClass(self.in_size)
             inputs = self.generate_input()
             state = neuron.init_state(self.batch_size)
-            call = bst.compile.jit(neuron)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(neuron)
+            with brainstate.environ.context(dt=0.1 * u.ms):
                 for t in range(self.time_steps):
                     out = call(inputs[t])
                     self.assertFalse(jax.tree_util.tree_leaves(out)[0].aval.weak_type)
@@ -148,15 +147,15 @@ class TestNeuron(unittest.TestCase):
             self.assertEqual(neuron.in_size, in_size)
             self.assertEqual(neuron.out_size, in_size)
 
-            inputs = bst.random.randn(self.time_steps, self.batch_size, *in_size) * u.mA
+            inputs = brainstate.random.randn(self.time_steps, self.batch_size, *in_size) * u.mA
             state = neuron.init_state(self.batch_size)
-            call = bst.compile.jit(neuron)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(neuron)
+            with brainstate.environ.context(dt=0.1 * u.ms):
                 for t in range(self.time_steps):
                     out = call(inputs[t])
                     self.assertEqual(out.shape, (self.batch_size, *in_size))
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/{_interaction/_normalizations.py → _normalizations.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import environ, init
 from brainstate._state import ParamState, BatchState
-from brainstate.nn._module import Module
 from brainstate.typing import DTypeLike, ArrayLike, Size, Axes
+from ._module import Module
 
 __all__ = [
     'BatchNorm0d',

brainstate/nn/{_interaction/_normalizations_test.py → _normalizations_test.py}

@@ -13,12 +13,10 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
-
 from absl.testing import absltest
 from absl.testing import parameterized
 
-import brainstate as bst
+import brainstate
 
 
 class Test_Normalization(parameterized.TestCase):
@@ -26,27 +24,27 @@ class Test_Normalization(parameterized.TestCase):
         fit=[True, False],
     )
     def test_BatchNorm1d(self, fit):
-        net = bst.nn.BatchNorm1d((3, 10))
-        bst.environ.set(fit=fit)
-        input = bst.random.randn(1, 3, 10)
+        net = brainstate.nn.BatchNorm1d((3, 10))
+        brainstate.environ.set(fit=fit)
+        input = brainstate.random.randn(1, 3, 10)
         output = net(input)
 
     @parameterized.product(
         fit=[True, False]
     )
     def test_BatchNorm2d(self, fit):
-        net = bst.nn.BatchNorm2d([3, 4, 10])
-        bst.environ.set(fit=fit)
-        input = bst.random.randn(1, 3, 4, 10)
+        net = brainstate.nn.BatchNorm2d([3, 4, 10])
+        brainstate.environ.set(fit=fit)
+        input = brainstate.random.randn(1, 3, 4, 10)
         output = net(input)
 
     @parameterized.product(
         fit=[True, False]
     )
     def test_BatchNorm3d(self, fit):
-        net = bst.nn.BatchNorm3d([3, 4, 5, 10])
-        bst.environ.set(fit=fit)
-        input = bst.random.randn(1, 3, 4, 5, 10)
+        net = brainstate.nn.BatchNorm3d([3, 4, 5, 10])
+        brainstate.environ.set(fit=fit)
+        input = brainstate.random.randn(1, 3, 4, 5, 10)
         output = net(input)
 
     # @parameterized.product(

brainstate/nn/{_interaction/_poolings.py → _poolings.py}

@@ -25,8 +25,8 @@ import jax.numpy as jnp
 import numpy as np
 
 from brainstate import environ
-from brainstate.nn._module import Module
 from brainstate.typing import Size
+from ._module import Module
 
 __all__ = [
     'Flatten', 'Unflatten',

brainstate/nn/{_interaction/_poolings_test.py → _poolings_test.py}

@@ -1,13 +1,11 @@
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import jax
 import numpy as np
 from absl.testing import absltest
 from absl.testing import parameterized
 
-import brainstate as bst
+import brainstate
 import brainstate.nn as nn
 
 
@@ -18,7 +16,7 @@ class TestFlatten(parameterized.TestCase):
             (32, 8),
             (10, 20, 30),
         ]:
-            arr = bst.random.rand(*size)
+            arr = brainstate.random.rand(*size)
             f = nn.Flatten(start_axis=0)
             out = f(arr)
             self.assertTrue(out.shape == (np.prod(size),))
@@ -29,21 +27,21 @@ class TestFlatten(parameterized.TestCase):
             (32, 8),
             (10, 20, 30),
         ]:
-            arr = bst.random.rand(*size)
+            arr = brainstate.random.rand(*size)
             f = nn.Flatten(start_axis=1)
             out = f(arr)
             self.assertTrue(out.shape == (size[0], np.prod(size[1:])))
 
     def test_flatten3(self):
         size = (16, 32, 32, 8)
-        arr = bst.random.rand(*size)
+        arr = brainstate.random.rand(*size)
         f = nn.Flatten(start_axis=0, in_size=(32, 8))
         out = f(arr)
         self.assertTrue(out.shape == (16, 32, 32 * 8))
 
     def test_flatten4(self):
         size = (16, 32, 32, 8)
-        arr = bst.random.rand(*size)
+        arr = brainstate.random.rand(*size)
         f = nn.Flatten(start_axis=1, in_size=(32, 32, 8))
         out = f(arr)
         self.assertTrue(out.shape == (16, 32, 32 * 8))
@@ -58,7 +56,7 @@ class TestPool(parameterized.TestCase):
         super().__init__(*args, **kwargs)
 
     def test_MaxPool2d_v1(self):
-        arr = bst.random.rand(16, 32, 32, 8)
+        arr = brainstate.random.rand(16, 32, 32, 8)
 
         out = nn.MaxPool2d(2, 2, channel_axis=-1)(arr)
         self.assertTrue(out.shape == (16, 16, 16, 8))
@@ -79,7 +77,7 @@ class TestPool(parameterized.TestCase):
         self.assertTrue(out.shape == (16, 17, 32, 5))
 
     def test_AvgPool2d_v1(self):
-        arr = bst.random.rand(16, 32, 32, 8)
+        arr = brainstate.random.rand(16, 32, 32, 8)
 
         out = nn.AvgPool2d(2, 2, channel_axis=-1)(arr)
         self.assertTrue(out.shape == (16, 16, 16, 8))
@@ -105,9 +103,9 @@ class TestPool(parameterized.TestCase):
         for target_size in [10, 9, 8, 7, 6]
     )
     def test_adaptive_pool1d(self, target_size):
-        from brainstate.nn._interaction._poolings import _adaptive_pool1d
+        from brainstate.nn._poolings import _adaptive_pool1d
 
-        arr = bst.random.rand(100)
+        arr = brainstate.random.rand(100)
         op = jax.numpy.mean
 
         out = _adaptive_pool1d(arr, target_size, op)
@@ -119,7 +117,7 @@ class TestPool(parameterized.TestCase):
         self.assertTrue(out.shape == (target_size,))
 
     def test_AdaptiveAvgPool2d_v1(self):
-        input = bst.random.randn(64, 8, 9)
+        input = brainstate.random.randn(64, 8, 9)
 
         output = nn.AdaptiveAvgPool2d((5, 7), channel_axis=0)(input)
         self.assertTrue(output.shape == (64, 5, 7))
@@ -137,8 +135,8 @@ class TestPool(parameterized.TestCase):
         self.assertTrue(output.shape == (64, 2, 3))
 
     def test_AdaptiveAvgPool2d_v2(self):
-        bst.random.seed()
-        input = bst.random.randn(128, 64, 32, 16)
+        brainstate.random.seed()
+        input = brainstate.random.randn(128, 64, 32, 16)
 
         output = nn.AdaptiveAvgPool2d((5, 7), channel_axis=0)(input)
         self.assertTrue(output.shape == (128, 64, 5, 7))
@@ -154,13 +152,13 @@ class TestPool(parameterized.TestCase):
         print()
 
     def test_AdaptiveAvgPool3d_v1(self):
-        input = bst.random.randn(10, 128, 64, 32)
+        input = brainstate.random.randn(10, 128, 64, 32)
         net = nn.AdaptiveAvgPool3d(target_size=[6, 5, 3], channel_axis=0)
         output = net(input)
         self.assertTrue(output.shape == (10, 6, 5, 3))
 
     def test_AdaptiveAvgPool3d_v2(self):
-        input = bst.random.randn(10, 20, 128, 64, 32)
+        input = brainstate.random.randn(10, 20, 128, 64, 32)
         net = nn.AdaptiveAvgPool3d(target_size=[6, 5, 3])
         output = net(input)
         self.assertTrue(output.shape == (10, 6, 5, 3, 32))
@@ -169,7 +167,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1)
     )
     def test_AdaptiveMaxPool1d_v1(self, axis):
-        input = bst.random.randn(32, 16)
+        input = brainstate.random.randn(32, 16)
         net = nn.AdaptiveMaxPool1d(target_size=4, channel_axis=axis)
         output = net(input)
 
@@ -177,7 +175,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1, 2)
     )
     def test_AdaptiveMaxPool1d_v2(self, axis):
-        input = bst.random.randn(2, 32, 16)
+        input = brainstate.random.randn(2, 32, 16)
         net = nn.AdaptiveMaxPool1d(target_size=4, channel_axis=axis)
         output = net(input)
 
@@ -185,7 +183,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1, 2)
     )
     def test_AdaptiveMaxPool2d_v1(self, axis):
-        input = bst.random.randn(32, 16, 12)
+        input = brainstate.random.randn(32, 16, 12)
         net = nn.AdaptiveAvgPool2d(target_size=[5, 4], channel_axis=axis)
         output = net(input)
 
@@ -193,7 +191,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1, 2, 3)
     )
     def test_AdaptiveMaxPool2d_v2(self, axis):
-        input = bst.random.randn(2, 32, 16, 12)
+        input = brainstate.random.randn(2, 32, 16, 12)
         net = nn.AdaptiveAvgPool2d(target_size=[5, 4], channel_axis=axis)
         output = net(input)
 
@@ -201,7 +199,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1, 2, 3)
     )
     def test_AdaptiveMaxPool3d_v1(self, axis):
-        input = bst.random.randn(2, 128, 64, 32)
+        input = brainstate.random.randn(2, 128, 64, 32)
         net = nn.AdaptiveMaxPool3d(target_size=[6, 5, 4], channel_axis=axis)
         output = net(input)
         print()
@@ -210,7 +208,7 @@ class TestPool(parameterized.TestCase):
         axis=(-1, 0, 1, 2, 3, 4)
     )
     def test_AdaptiveMaxPool3d_v1(self, axis):
-        input = bst.random.randn(2, 128, 64, 32, 16)
+        input = brainstate.random.randn(2, 128, 64, 32, 16)
         net = nn.AdaptiveMaxPool3d(target_size=[6, 5, 4], channel_axis=axis)
         output = net(input)