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

brainstate/nn/_elementwise/_dropout_test.py

@@ -18,19 +18,19 @@ import unittest
 
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestDropout(unittest.TestCase):
 
     def test_dropout(self):
         # Create a Dropout layer with a dropout rate of 0.5
-        dropout_layer = bst.nn.Dropout(0.5)
+        dropout_layer = brainstate.nn.Dropout(0.5)
 
         # Input data
         input_data = np.arange(20)
 
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             # Apply dropout
             output_data = dropout_layer(input_data)
 
@@ -47,10 +47,10 @@ class TestDropout(unittest.TestCase):
             np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements)
 
     def test_DropoutFixed(self):
-        dropout_layer = bst.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
+        dropout_layer = brainstate.nn.DropoutFixed(in_size=(2, 3), prob=0.5)
         dropout_layer.init_state(batch_size=2)
         input_data = np.random.randn(2, 2, 3)
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             output_data = dropout_layer.update(input_data)
         self.assertEqual(input_data.shape, output_data.shape)
         self.assertTrue(np.any(output_data == 0))
@@ -72,9 +72,9 @@ class TestDropout(unittest.TestCase):
     #     np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=4)
 
     def test_Dropout2d(self):
-        dropout_layer = bst.nn.Dropout2d(prob=0.5)
+        dropout_layer = brainstate.nn.Dropout2d(prob=0.5)
         input_data = np.random.randn(2, 3, 4, 5)
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             output_data = dropout_layer(input_data)
         self.assertEqual(input_data.shape, output_data.shape)
         self.assertTrue(np.any(output_data == 0))
@@ -84,9 +84,9 @@ class TestDropout(unittest.TestCase):
         np.testing.assert_almost_equal(non_zero_elements, expected_non_zero_elements, decimal=4)
 
     def test_Dropout3d(self):
-        dropout_layer = bst.nn.Dropout3d(prob=0.5)
+        dropout_layer = brainstate.nn.Dropout3d(prob=0.5)
         input_data = np.random.randn(2, 3, 4, 5, 6)
-        with bst.environ.context(fit=True):
+        with brainstate.environ.context(fit=True):
             output_data = dropout_layer(input_data)
         self.assertEqual(input_data.shape, output_data.shape)
         self.assertTrue(np.any(output_data == 0))

brainstate/nn/_elementwise/_elementwise_test.py

@@ -13,157 +13,155 @@
 # 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_Activation(parameterized.TestCase):
 
     def test_Threshold(self):
-        threshold_layer = bst.nn.Threshold(5, 20)
-        input = bst.random.randn(2)
+        threshold_layer = brainstate.nn.Threshold(5, 20)
+        input = brainstate.random.randn(2)
         output = threshold_layer(input)
 
     def test_ReLU(self):
-        ReLU_layer = bst.nn.ReLU()
-        input = bst.random.randn(2)
+        ReLU_layer = brainstate.nn.ReLU()
+        input = brainstate.random.randn(2)
         output = ReLU_layer(input)
 
     def test_RReLU(self):
-        RReLU_layer = bst.nn.RReLU(lower=0, upper=1)
-        input = bst.random.randn(2)
+        RReLU_layer = brainstate.nn.RReLU(lower=0, upper=1)
+        input = brainstate.random.randn(2)
         output = RReLU_layer(input)
 
     def test_Hardtanh(self):
-        Hardtanh_layer = bst.nn.Hardtanh(min_val=0, max_val=1, )
-        input = bst.random.randn(2)
+        Hardtanh_layer = brainstate.nn.Hardtanh(min_val=0, max_val=1, )
+        input = brainstate.random.randn(2)
         output = Hardtanh_layer(input)
 
     def test_ReLU6(self):
-        ReLU6_layer = bst.nn.ReLU6()
-        input = bst.random.randn(2)
+        ReLU6_layer = brainstate.nn.ReLU6()
+        input = brainstate.random.randn(2)
         output = ReLU6_layer(input)
 
     def test_Sigmoid(self):
-        Sigmoid_layer = bst.nn.Sigmoid()
-        input = bst.random.randn(2)
+        Sigmoid_layer = brainstate.nn.Sigmoid()
+        input = brainstate.random.randn(2)
         output = Sigmoid_layer(input)
 
     def test_Hardsigmoid(self):
-        Hardsigmoid_layer = bst.nn.Hardsigmoid()
-        input = bst.random.randn(2)
+        Hardsigmoid_layer = brainstate.nn.Hardsigmoid()
+        input = brainstate.random.randn(2)
         output = Hardsigmoid_layer(input)
 
     def test_Tanh(self):
-        Tanh_layer = bst.nn.Tanh()
-        input = bst.random.randn(2)
+        Tanh_layer = brainstate.nn.Tanh()
+        input = brainstate.random.randn(2)
         output = Tanh_layer(input)
 
     def test_SiLU(self):
-        SiLU_layer = bst.nn.SiLU()
-        input = bst.random.randn(2)
+        SiLU_layer = brainstate.nn.SiLU()
+        input = brainstate.random.randn(2)
         output = SiLU_layer(input)
 
     def test_Mish(self):
-        Mish_layer = bst.nn.Mish()
-        input = bst.random.randn(2)
+        Mish_layer = brainstate.nn.Mish()
+        input = brainstate.random.randn(2)
         output = Mish_layer(input)
 
     def test_Hardswish(self):
-        Hardswish_layer = bst.nn.Hardswish()
-        input = bst.random.randn(2)
+        Hardswish_layer = brainstate.nn.Hardswish()
+        input = brainstate.random.randn(2)
         output = Hardswish_layer(input)
 
     def test_ELU(self):
-        ELU_layer = bst.nn.ELU(alpha=0.5, )
-        input = bst.random.randn(2)
+        ELU_layer = brainstate.nn.ELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
         output = ELU_layer(input)
 
     def test_CELU(self):
-        CELU_layer = bst.nn.CELU(alpha=0.5, )
-        input = bst.random.randn(2)
+        CELU_layer = brainstate.nn.CELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
         output = CELU_layer(input)
 
     def test_SELU(self):
-        SELU_layer = bst.nn.SELU()
-        input = bst.random.randn(2)
+        SELU_layer = brainstate.nn.SELU()
+        input = brainstate.random.randn(2)
         output = SELU_layer(input)
 
     def test_GLU(self):
-        GLU_layer = bst.nn.GLU()
-        input = bst.random.randn(4, 2)
+        GLU_layer = brainstate.nn.GLU()
+        input = brainstate.random.randn(4, 2)
         output = GLU_layer(input)
 
     @parameterized.product(
         approximate=['tanh', 'none']
     )
     def test_GELU(self, approximate):
-        GELU_layer = bst.nn.GELU()
-        input = bst.random.randn(2)
+        GELU_layer = brainstate.nn.GELU()
+        input = brainstate.random.randn(2)
         output = GELU_layer(input)
 
     def test_Hardshrink(self):
-        Hardshrink_layer = bst.nn.Hardshrink(lambd=1)
-        input = bst.random.randn(2)
+        Hardshrink_layer = brainstate.nn.Hardshrink(lambd=1)
+        input = brainstate.random.randn(2)
         output = Hardshrink_layer(input)
 
     def test_LeakyReLU(self):
-        LeakyReLU_layer = bst.nn.LeakyReLU()
-        input = bst.random.randn(2)
+        LeakyReLU_layer = brainstate.nn.LeakyReLU()
+        input = brainstate.random.randn(2)
         output = LeakyReLU_layer(input)
 
     def test_LogSigmoid(self):
-        LogSigmoid_layer = bst.nn.LogSigmoid()
-        input = bst.random.randn(2)
+        LogSigmoid_layer = brainstate.nn.LogSigmoid()
+        input = brainstate.random.randn(2)
         output = LogSigmoid_layer(input)
 
     def test_Softplus(self):
-        Softplus_layer = bst.nn.Softplus()
-        input = bst.random.randn(2)
+        Softplus_layer = brainstate.nn.Softplus()
+        input = brainstate.random.randn(2)
         output = Softplus_layer(input)
 
     def test_Softshrink(self):
-        Softshrink_layer = bst.nn.Softshrink(lambd=1)
-        input = bst.random.randn(2)
+        Softshrink_layer = brainstate.nn.Softshrink(lambd=1)
+        input = brainstate.random.randn(2)
         output = Softshrink_layer(input)
 
     def test_PReLU(self):
-        PReLU_layer = bst.nn.PReLU(num_parameters=2, init=0.5)
-        input = bst.random.randn(2)
+        PReLU_layer = brainstate.nn.PReLU(num_parameters=2, init=0.5)
+        input = brainstate.random.randn(2)
         output = PReLU_layer(input)
 
     def test_Softsign(self):
-        Softsign_layer = bst.nn.Softsign()
-        input = bst.random.randn(2)
+        Softsign_layer = brainstate.nn.Softsign()
+        input = brainstate.random.randn(2)
         output = Softsign_layer(input)
 
     def test_Tanhshrink(self):
-        Tanhshrink_layer = bst.nn.Tanhshrink()
-        input = bst.random.randn(2)
+        Tanhshrink_layer = brainstate.nn.Tanhshrink()
+        input = brainstate.random.randn(2)
         output = Tanhshrink_layer(input)
 
     def test_Softmin(self):
-        Softmin_layer = bst.nn.Softmin(dim=2)
-        input = bst.random.randn(2, 3, 4)
+        Softmin_layer = brainstate.nn.Softmin(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
         output = Softmin_layer(input)
 
     def test_Softmax(self):
-        Softmax_layer = bst.nn.Softmax(dim=2)
-        input = bst.random.randn(2, 3, 4)
+        Softmax_layer = brainstate.nn.Softmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
         output = Softmax_layer(input)
 
     def test_Softmax2d(self):
-        Softmax2d_layer = bst.nn.Softmax2d()
-        input = bst.random.randn(2, 3, 12, 13)
+        Softmax2d_layer = brainstate.nn.Softmax2d()
+        input = brainstate.random.randn(2, 3, 12, 13)
         output = Softmax2d_layer(input)
 
     def test_LogSoftmax(self):
-        LogSoftmax_layer = bst.nn.LogSoftmax(dim=2)
-        input = bst.random.randn(2, 3, 4)
+        LogSoftmax_layer = brainstate.nn.LogSoftmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
         output = LogSoftmax_layer(input)
 
 

brainstate/nn/_event/_fixedprob_mv_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import jax.numpy
 import jax.numpy as jnp

brainstate/nn/_event/_linear_mv_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import jax
 import jax.numpy as jnp

brainstate/nn/_exp_euler_test.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import brainunit as u
 
-import brainstate as bst
+import brainstate
 
 
 class TestExpEuler(unittest.TestCase):
@@ -27,10 +26,10 @@ class TestExpEuler(unittest.TestCase):
         def fun(x, tau):
             return -x / tau
 
-        with bst.environ.context(dt=0.1):
+        with brainstate.environ.context(dt=0.1):
             with self.assertRaises(AssertionError):
-                r = bst.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
+                r = brainstate.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
 
-        with bst.environ.context(dt=1. * u.ms):
-            r = bst.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
+        with brainstate.environ.context(dt=1. * u.ms):
+            r = brainstate.nn.exp_euler_step(fun, 1.0 * u.mV, 1. * u.ms)
             print(r)

brainstate/nn/_interaction/_conv_test.py

@@ -1,13 +1,11 @@
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import jax.numpy as jnp
 import pytest
 from absl.testing import absltest
 from absl.testing import parameterized
 
-import brainstate as bst
+import brainstate
 
 
 class TestConv(parameterized.TestCase):
@@ -19,8 +17,8 @@ class TestConv(parameterized.TestCase):
             img = img.at[0, x:x + 10, y:y + 10, k].set(1.0)
             img = img.at[1, x:x + 20, y:y + 20, k].set(3.0)
 
-        net = bst.nn.Conv2d((200, 198, 4), out_channels=32, kernel_size=(3, 3),
-                            stride=(2, 1), padding='VALID', groups=4)
+        net = brainstate.nn.Conv2d((200, 198, 4), out_channels=32, kernel_size=(3, 3),
+                                   stride=(2, 1), padding='VALID', groups=4)
         out = net(img)
         print("out shape: ", out.shape)
         self.assertEqual(out.shape, (2, 99, 196, 32))
@@ -30,7 +28,7 @@ class TestConv(parameterized.TestCase):
         # plt.show()
 
     def test_conv1D(self):
-        model = bst.nn.Conv1d((5, 3), out_channels=32, kernel_size=(3,))
+        model = brainstate.nn.Conv1d((5, 3), out_channels=32, kernel_size=(3,))
         input = jnp.ones((2, 5, 3))
         out = model(input)
         print("out shape: ", out.shape)
@@ -41,7 +39,7 @@ class TestConv(parameterized.TestCase):
         # plt.show()
 
     def test_conv2D(self):
-        model = bst.nn.Conv2d((5, 5, 3), out_channels=32, kernel_size=(3, 3))
+        model = brainstate.nn.Conv2d((5, 5, 3), out_channels=32, kernel_size=(3, 3))
         input = jnp.ones((2, 5, 5, 3))
 
         out = model(input)
@@ -49,7 +47,7 @@ class TestConv(parameterized.TestCase):
         self.assertEqual(out.shape, (2, 5, 5, 32))
 
     def test_conv3D(self):
-        model = bst.nn.Conv3d((5, 5, 5, 3), out_channels=32, kernel_size=(3, 3, 3))
+        model = brainstate.nn.Conv3d((5, 5, 5, 3), out_channels=32, kernel_size=(3, 3, 3))
         input = jnp.ones((2, 5, 5, 5, 3))
         out = model(input)
         print("out shape: ", out.shape)
@@ -62,13 +60,13 @@ class TestConvTranspose1d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 3))
         for use_bias in [True, False]:
-            conv_transpose_module = bst.nn.ConvTranspose1d(
+            conv_transpose_module = brainstate.nn.ConvTranspose1d(
                 in_channels=3,
                 out_channels=4,
                 kernel_size=(3,),
                 padding='VALID',
-                w_initializer=bst.init.Constant(1.),
-                b_initializer=bst.init.Constant(1.) if use_bias else None,
+                w_initializer=brainstate.init.Constant(1.),
+                b_initializer=brainstate.init.Constant(1.) if use_bias else None,
             )
             self.assertEqual(conv_transpose_module.w.shape, (3, 3, 4))
             y = conv_transpose_module(x)
@@ -91,14 +89,14 @@ class TestConvTranspose1d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 3))
         m = jnp.tril(jnp.ones((3, 3, 4)))
-        conv_transpose_module = bst.nn.ConvTranspose1d(
+        conv_transpose_module = brainstate.nn.ConvTranspose1d(
             in_channels=3,
             out_channels=4,
             kernel_size=(3,),
             padding='VALID',
             mask=m,
-            w_initializer=bst.init.Constant(),
-            b_initializer=bst.init.Constant(),
+            w_initializer=brainstate.init.Constant(),
+            b_initializer=brainstate.init.Constant(),
         )
         self.assertEqual(conv_transpose_module.w.shape, (3, 3, 4))
         y = conv_transpose_module(x)
@@ -119,14 +117,14 @@ class TestConvTranspose1d(parameterized.TestCase):
 
         data = jnp.ones([1, 3, 1])
         for use_bias in [True, False]:
-            net = bst.nn.ConvTranspose1d(
+            net = brainstate.nn.ConvTranspose1d(
                 in_channels=1,
                 out_channels=1,
                 kernel_size=3,
                 stride=1,
                 padding="SAME",
-                w_initializer=bst.init.Constant(),
-                b_initializer=bst.init.Constant() if use_bias else None,
+                w_initializer=brainstate.init.Constant(),
+                b_initializer=brainstate.init.Constant() if use_bias else None,
             )
             out = net(data)
             self.assertEqual(out.shape, (1, 3, 1))
@@ -143,13 +141,13 @@ class TestConvTranspose2d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 3))
         for use_bias in [True, False]:
-            conv_transpose_module = bst.nn.ConvTranspose2d(
+            conv_transpose_module = brainstate.nn.ConvTranspose2d(
                 in_channels=3,
                 out_channels=4,
                 kernel_size=(3, 3),
                 padding='VALID',
-                w_initializer=bst.init.Constant(),
-                b_initializer=bst.init.Constant() if use_bias else None,
+                w_initializer=brainstate.init.Constant(),
+                b_initializer=brainstate.init.Constant() if use_bias else None,
             )
         self.assertEqual(conv_transpose_module.w.shape, (3, 3, 3, 4))
         y = conv_transpose_module(x)
@@ -159,13 +157,13 @@ class TestConvTranspose2d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 3))
         m = jnp.tril(jnp.ones((3, 3, 3, 4)))
-        conv_transpose_module = bst.nn.ConvTranspose2d(
+        conv_transpose_module = brainstate.nn.ConvTranspose2d(
             in_channels=3,
             out_channels=4,
             kernel_size=(3, 3),
             padding='VALID',
             mask=m,
-            w_initializer=bst.init.Constant(),
+            w_initializer=brainstate.init.Constant(),
         )
         y = conv_transpose_module(x)
         print(y.shape)
@@ -174,14 +172,14 @@ class TestConvTranspose2d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 3))
         for use_bias in [True, False]:
-            conv_transpose_module = bst.nn.ConvTranspose2d(
+            conv_transpose_module = brainstate.nn.ConvTranspose2d(
                 in_channels=3,
                 out_channels=4,
                 kernel_size=(3, 3),
                 stride=1,
                 padding='SAME',
-                w_initializer=bst.init.Constant(),
-                b_initializer=bst.init.Constant() if use_bias else None,
+                w_initializer=brainstate.init.Constant(),
+                b_initializer=brainstate.init.Constant() if use_bias else None,
             )
         y = conv_transpose_module(x)
         print(y.shape)
@@ -193,13 +191,13 @@ class TestConvTranspose3d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 8, 3))
         for use_bias in [True, False]:
-            conv_transpose_module = bst.nn.ConvTranspose3d(
+            conv_transpose_module = brainstate.nn.ConvTranspose3d(
                 in_channels=3,
                 out_channels=4,
                 kernel_size=(3, 3, 3),
                 padding='VALID',
-                w_initializer=bst.init.Constant(),
-                b_initializer=bst.init.Constant() if use_bias else None,
+                w_initializer=brainstate.init.Constant(),
+                b_initializer=brainstate.init.Constant() if use_bias else None,
             )
         y = conv_transpose_module(x)
         print(y.shape)
@@ -208,13 +206,13 @@ class TestConvTranspose3d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 8, 3))
         m = jnp.tril(jnp.ones((3, 3, 3, 3, 4)))
-        conv_transpose_module = bst.nn.ConvTranspose3d(
+        conv_transpose_module = brainstate.nn.ConvTranspose3d(
             in_channels=3,
             out_channels=4,
             kernel_size=(3, 3, 3),
             padding='VALID',
             mask=m,
-            w_initializer=bst.init.Constant(),
+            w_initializer=brainstate.init.Constant(),
         )
         y = conv_transpose_module(x)
         print(y.shape)
@@ -223,14 +221,14 @@ class TestConvTranspose3d(parameterized.TestCase):
 
         x = jnp.ones((1, 8, 8, 8, 3))
         for use_bias in [True, False]:
-            conv_transpose_module = bst.nn.ConvTranspose3d(
+            conv_transpose_module = brainstate.nn.ConvTranspose3d(
                 in_channels=3,
                 out_channels=4,
                 kernel_size=(3, 3, 3),
                 stride=1,
                 padding='SAME',
-                w_initializer=bst.init.Constant(),
-                b_initializer=bst.init.Constant() if use_bias else None,
+                w_initializer=brainstate.init.Constant(),
+                b_initializer=brainstate.init.Constant() if use_bias else None,
             )
         y = conv_transpose_module(x)
         print(y.shape)

brainstate/nn/_interaction/_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/_interaction/_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(