brainstate 0.1.0.post20250503__py2.py3-none-any.whl → 0.1.2__py2.py3-none-any.whl

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/_embedding.py

@@ -12,7 +12,6 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # ==============================================================================
-from __future__ import annotations
 
 from typing import Optional, Callable, Union
 

brainstate/nn/_interaction/_linear.py

@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 from typing import Callable, Union, Optional
 
 import brainunit as u

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.py

@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 from typing import Callable, Union, Sequence, Optional, Any
 
 import jax

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(

brainstate/nn/_interaction/_poolings.py

@@ -15,8 +15,6 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
-
 import functools
 from typing import Sequence, Optional
 from typing import Union, Tuple, Callable, List

brainstate/nn/_interaction/_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))
@@ -107,7 +105,7 @@ class TestPool(parameterized.TestCase):
     def test_adaptive_pool1d(self, target_size):
         from brainstate.nn._interaction._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)
 

brainstate/nn/_module.py

@@ -25,7 +25,6 @@ The basic classes include:
 - ``Sequential``: The class for a sequential of modules, which update the modules sequentially.
 
 """
-from __future__ import annotations
 
 import warnings
 from typing import Sequence, Optional, Tuple, Union, TYPE_CHECKING, Callable