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

brainstate/init/_random_inits_test.py

@@ -14,30 +14,29 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 import unittest
 
-import brainstate as bst
+import brainstate
 
 
 class TestNormalInit(unittest.TestCase):
 
     def test_normal_init1(self):
-        init = bst.init.Normal()
+        init = brainstate.init.Normal()
         for size in [(100,), (10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_normal_init2(self):
-        init = bst.init.Normal(scale=0.5)
+        init = brainstate.init.Normal(scale=0.5)
         for size in [(100,), (10, 20)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_normal_init3(self):
-        init1 = bst.init.Normal(scale=0.5, seed=10)
-        init2 = bst.init.Normal(scale=0.5, seed=10)
+        init1 = brainstate.init.Normal(scale=0.5, seed=10)
+        init2 = brainstate.init.Normal(scale=0.5, seed=10)
         size = (10,)
         weights1 = init1(size)
         weights2 = init2(size)
@@ -47,13 +46,13 @@ class TestNormalInit(unittest.TestCase):
 
 class TestUniformInit(unittest.TestCase):
     def test_uniform_init1(self):
-        init = bst.init.Normal()
+        init = brainstate.init.Normal()
         for size in [(100,), (10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_uniform_init2(self):
-        init = bst.init.Uniform(min_val=10, max_val=20)
+        init = brainstate.init.Uniform(min_val=10, max_val=20)
         for size in [(100,), (10, 20)]:
             weights = init(size)
             assert weights.shape == size
@@ -61,20 +60,20 @@ class TestUniformInit(unittest.TestCase):
 
 class TestVarianceScaling(unittest.TestCase):
     def test_var_scaling1(self):
-        init = bst.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
+        init = brainstate.init.VarianceScaling(scale=1., mode='fan_in', distribution='truncated_normal')
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_var_scaling2(self):
-        init = bst.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
+        init = brainstate.init.VarianceScaling(scale=2, mode='fan_out', distribution='normal')
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_var_scaling3(self):
-        init = bst.init.VarianceScaling(scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1,
-                                        distribution='uniform')
+        init = brainstate.init.VarianceScaling(scale=2 / 4, mode='fan_avg', in_axis=0, out_axis=1,
+                                               distribution='uniform')
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -82,7 +81,7 @@ class TestVarianceScaling(unittest.TestCase):
 
 class TestKaimingUniformUnit(unittest.TestCase):
     def test_kaiming_uniform_init(self):
-        init = bst.init.KaimingUniform()
+        init = brainstate.init.KaimingUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -90,7 +89,7 @@ class TestKaimingUniformUnit(unittest.TestCase):
 
 class TestKaimingNormalUnit(unittest.TestCase):
     def test_kaiming_normal_init(self):
-        init = bst.init.KaimingNormal()
+        init = brainstate.init.KaimingNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -98,7 +97,7 @@ class TestKaimingNormalUnit(unittest.TestCase):
 
 class TestXavierUniformUnit(unittest.TestCase):
     def test_xavier_uniform_init(self):
-        init = bst.init.XavierUniform()
+        init = brainstate.init.XavierUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -106,7 +105,7 @@ class TestXavierUniformUnit(unittest.TestCase):
 
 class TestXavierNormalUnit(unittest.TestCase):
     def test_xavier_normal_init(self):
-        init = bst.init.XavierNormal()
+        init = brainstate.init.XavierNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -114,7 +113,7 @@ class TestXavierNormalUnit(unittest.TestCase):
 
 class TestLecunUniformUnit(unittest.TestCase):
     def test_lecun_uniform_init(self):
-        init = bst.init.LecunUniform()
+        init = brainstate.init.LecunUniform()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -122,7 +121,7 @@ class TestLecunUniformUnit(unittest.TestCase):
 
 class TestLecunNormalUnit(unittest.TestCase):
     def test_lecun_normal_init(self):
-        init = bst.init.LecunNormal()
+        init = brainstate.init.LecunNormal()
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -130,13 +129,13 @@ class TestLecunNormalUnit(unittest.TestCase):
 
 class TestOrthogonalUnit(unittest.TestCase):
     def test_orthogonal_init1(self):
-        init = bst.init.Orthogonal()
+        init = brainstate.init.Orthogonal()
         for size in [(20, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
 
     def test_orthogonal_init2(self):
-        init = bst.init.Orthogonal(scale=2., axis=0)
+        init = brainstate.init.Orthogonal(scale=2., axis=0)
         for size in [(10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -144,7 +143,7 @@ class TestOrthogonalUnit(unittest.TestCase):
 
 class TestDeltaOrthogonalUnit(unittest.TestCase):
     def test_delta_orthogonal_init1(self):
-        init = bst.init.DeltaOrthogonal()
+        init = brainstate.init.DeltaOrthogonal()
         for size in [(20, 20, 20), (10, 20, 30, 40), (50, 40, 30, 20, 20)]:
             weights = init(size)
             assert weights.shape == size

brainstate/init/_regular_inits_test.py

@@ -14,16 +14,15 @@
 # ==============================================================================
 
 # -*- coding: utf-8 -*-
-from __future__ import annotations
 
 import unittest
 
-import brainstate as bst
+import brainstate
 
 
 class TestZeroInit(unittest.TestCase):
     def test_zero_init(self):
-        init = bst.init.ZeroInit()
+        init = brainstate.init.ZeroInit()
         for size in [(100,), (10, 20), (10, 20, 30)]:
             weights = init(size)
             assert weights.shape == size
@@ -33,7 +32,7 @@ class TestOneInit(unittest.TestCase):
     def test_one_init(self):
         for size in [(100,), (10, 20), (10, 20, 30)]:
             for value in [0., 1., -1.]:
-                init = bst.init.Constant(value=value)
+                init = brainstate.init.Constant(value=value)
                 weights = init(size)
                 assert weights.shape == size
                 assert (weights == value).all()
@@ -43,7 +42,7 @@ class TestIdentityInit(unittest.TestCase):
     def test_identity_init(self):
         for size in [(100,), (10, 20)]:
             for value in [0., 1., -1.]:
-                init = bst.init.Identity(value=value)
+                init = brainstate.init.Identity(value=value)
                 weights = init(size)
                 if len(size) == 1:
                     assert weights.shape == (size[0], size[0])

brainstate/nn/_collective_ops_test.py

@@ -16,21 +16,21 @@
 # -*- coding: utf-8 -*-
 
 
-import brainstate as bst
+import brainstate
 
 
 class Test_vmap_init_all_states:
 
     def test_vmap_init_all_states(self):
-        gru = bst.nn.GRUCell(1, 2)
-        bst.nn.vmap_init_all_states(gru, axis_size=10)
+        gru = brainstate.nn.GRUCell(1, 2)
+        brainstate.nn.vmap_init_all_states(gru, axis_size=10)
         print(gru)
 
     def test_vmap_init_all_states_v2(self):
-        @bst.compile.jit
+        @brainstate.compile.jit
         def init():
-            gru = bst.nn.GRUCell(1, 2)
-            bst.nn.vmap_init_all_states(gru, axis_size=10)
+            gru = brainstate.nn.GRUCell(1, 2)
+            brainstate.nn.vmap_init_all_states(gru, axis_size=10)
             print(gru)
 
         init()
@@ -38,6 +38,6 @@ class Test_vmap_init_all_states:
 
 class Test_init_all_states:
     def test_init_all_states(self):
-        gru = bst.nn.GRUCell(1, 2)
-        bst.nn.init_all_states(gru, batch_size=10)
+        gru = brainstate.nn.GRUCell(1, 2)
+        brainstate.nn.init_all_states(gru, batch_size=10)
         print(gru)

brainstate/nn/_dyn_impl/_dynamics_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/_dyn_impl/_dynamics_synapse_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -21,7 +20,7 @@ import brainunit as u
 import jax.numpy as jnp
 import pytest
 
-import brainstate as bst
+import brainstate
 from brainstate.nn import Expon, STP, STD
 
 
@@ -32,7 +31,7 @@ class TestSynapse(unittest.TestCase):
         self.time_steps = 100
 
     def generate_input(self):
-        return bst.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mS
+        return brainstate.random.randn(self.time_steps, self.batch_size, self.in_size) * u.mS
 
     def test_expon_synapse(self):
         tau = 20.0 * u.ms
@@ -46,8 +45,8 @@ class TestSynapse(unittest.TestCase):
 
         # Test forward pass
         state = synapse.init_state(self.batch_size)
-        call = bst.compile.jit(synapse)
-        with bst.environ.context(dt=0.1 * u.ms):
+        call = brainstate.compile.jit(synapse)
+        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))
@@ -75,7 +74,7 @@ class TestSynapse(unittest.TestCase):
 
         # Test forward pass
         state = synapse.init_state(self.batch_size)
-        call = bst.compile.jit(synapse)
+        call = brainstate.compile.jit(synapse)
         for t in range(self.time_steps):
             out = call(inputs[t])
             self.assertEqual(out.shape, (self.batch_size, self.in_size))
@@ -118,15 +117,15 @@ class TestSynapse(unittest.TestCase):
             self.assertEqual(synapse.in_size, in_size)
             self.assertEqual(synapse.out_size, in_size)
 
-            inputs = bst.random.randn(self.time_steps, self.batch_size, *in_size) * u.mS
+            inputs = brainstate.random.randn(self.time_steps, self.batch_size, *in_size) * u.mS
             state = synapse.init_state(self.batch_size)
-            call = bst.compile.jit(synapse)
-            with bst.environ.context(dt=0.1 * u.ms):
+            call = brainstate.compile.jit(synapse)
+            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/_dyn_impl/_rate_rnns_test.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestRateRNNModels(unittest.TestCase):
@@ -30,31 +29,31 @@ class TestRateRNNModels(unittest.TestCase):
         self.x = jnp.ones((self.batch_size, self.num_in))
 
     def test_ValinaRNNCell(self):
-        model = bst.nn.ValinaRNNCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.ValinaRNNCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_GRUCell(self):
-        model = bst.nn.GRUCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.GRUCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_MGUCell(self):
-        model = bst.nn.MGUCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.MGUCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_LSTMCell(self):
-        model = bst.nn.LSTMCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.LSTMCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))
 
     def test_URLSTMCell(self):
-        model = bst.nn.URLSTMCell(num_in=self.num_in, num_out=self.num_out)
+        model = brainstate.nn.URLSTMCell(num_in=self.num_in, num_out=self.num_out)
         model.init_state(batch_size=self.batch_size)
         output = model.update(self.x)
         self.assertEqual(output.shape, (self.batch_size, self.num_out))

brainstate/nn/_dyn_impl/_readout_test.py

@@ -13,13 +13,12 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
 import jax.numpy as jnp
 
-import brainstate as bst
+import brainstate
 
 
 class TestReadoutModels(unittest.TestCase):
@@ -32,23 +31,23 @@ class TestReadoutModels(unittest.TestCase):
         self.x = jnp.ones((self.batch_size, self.in_size))
 
     def test_LeakyRateReadout(self):
-        with bst.environ.context(dt=0.1):
-            model = bst.nn.LeakyRateReadout(in_size=self.in_size, out_size=self.out_size, tau=self.tau)
+        with brainstate.environ.context(dt=0.1):
+            model = brainstate.nn.LeakyRateReadout(in_size=self.in_size, out_size=self.out_size, tau=self.tau)
             model.init_state(batch_size=self.batch_size)
             output = model.update(self.x)
             self.assertEqual(output.shape, (self.batch_size, self.out_size))
 
     def test_LeakySpikeReadout(self):
-        with bst.environ.context(dt=0.1):
-            model = bst.nn.LeakySpikeReadout(in_size=self.in_size, tau=self.tau, V_th=self.V_th,
-                                             V_initializer=bst.init.ZeroInit(),
-                                             w_init=bst.init.KaimingNormal())
+        with brainstate.environ.context(dt=0.1):
+            model = brainstate.nn.LeakySpikeReadout(in_size=self.in_size, tau=self.tau, V_th=self.V_th,
+                                                    V_initializer=brainstate.init.ZeroInit(),
+                                                    w_init=brainstate.init.KaimingNormal())
             model.init_state(batch_size=self.batch_size)
-            with bst.environ.context(t=0.):
+            with brainstate.environ.context(t=0.):
                 output = model.update(self.x)
             self.assertEqual(output.shape, (self.batch_size, self.out_size))
 
 
 if __name__ == '__main__':
-    with bst.environ.context(dt=0.1):
+    with brainstate.environ.context(dt=0.1):
         unittest.main()

brainstate/nn/_dynamics/_dynamics_base_test.py

@@ -15,47 +15,46 @@
 
 # -*- coding: utf-8 -*-
 
-from __future__ import annotations
 
 import unittest
 
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestModuleGroup(unittest.TestCase):
     def test_initialization(self):
-        group = bst.nn.DynamicsGroup()
-        self.assertIsInstance(group, bst.nn.DynamicsGroup)
+        group = brainstate.nn.DynamicsGroup()
+        self.assertIsInstance(group, brainstate.nn.DynamicsGroup)
 
 
 class TestProjection(unittest.TestCase):
     def test_initialization(self):
-        proj = bst.nn.Projection()
-        self.assertIsInstance(proj, bst.nn.Projection)
+        proj = brainstate.nn.Projection()
+        self.assertIsInstance(proj, brainstate.nn.Projection)
 
     def test_update_not_implemented(self):
-        proj = bst.nn.Projection()
+        proj = brainstate.nn.Projection()
         with self.assertRaises(ValueError):
             proj.update()
 
 
 class TestDynamics(unittest.TestCase):
     def test_initialization(self):
-        dyn = bst.nn.Dynamics(in_size=10)
-        self.assertIsInstance(dyn, bst.nn.Dynamics)
+        dyn = brainstate.nn.Dynamics(in_size=10)
+        self.assertIsInstance(dyn, brainstate.nn.Dynamics)
         self.assertEqual(dyn.in_size, (10,))
         self.assertEqual(dyn.out_size, (10,))
 
     def test_size_validation(self):
         with self.assertRaises(ValueError):
-            bst.nn.Dynamics(in_size=[])
+            brainstate.nn.Dynamics(in_size=[])
         with self.assertRaises(ValueError):
-            bst.nn.Dynamics(in_size="invalid")
+            brainstate.nn.Dynamics(in_size="invalid")
 
     def test_input_handling(self):
-        dyn = bst.nn.Dynamics(in_size=10)
+        dyn = brainstate.nn.Dynamics(in_size=10)
         dyn.add_current_input("test_current", lambda: np.random.rand(10))
         dyn.add_delta_input("test_delta", lambda: np.random.rand(10))
 
@@ -63,15 +62,15 @@ class TestDynamics(unittest.TestCase):
         self.assertIn("test_delta", dyn.delta_inputs)
 
     def test_duplicate_input_key(self):
-        dyn = bst.nn.Dynamics(in_size=10)
+        dyn = brainstate.nn.Dynamics(in_size=10)
         dyn.add_current_input("test", lambda: np.random.rand(10))
         with self.assertRaises(ValueError):
             dyn.add_current_input("test", lambda: np.random.rand(10))
 
     def test_varshape(self):
-        dyn = bst.nn.Dynamics(in_size=(2, 3))
+        dyn = brainstate.nn.Dynamics(in_size=(2, 3))
         self.assertEqual(dyn.varshape, (2, 3))
-        dyn = bst.nn.Dynamics(in_size=(2, 3))
+        dyn = brainstate.nn.Dynamics(in_size=(2, 3))
         self.assertEqual(dyn.varshape, (2, 3))
 
 

brainstate/nn/_dynamics/_synouts_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -21,7 +20,7 @@ import brainunit as u
 import jax.numpy as jnp
 import numpy as np
 
-import brainstate as bst
+import brainstate
 
 
 class TestSynOutModels(unittest.TestCase):
@@ -35,19 +34,19 @@ class TestSynOutModels(unittest.TestCase):
         self.V_offset = jnp.array([0.0])
 
     def test_COBA(self):
-        model = bst.nn.COBA(E=self.E)
+        model = brainstate.nn.COBA(E=self.E)
         output = model.update(self.conductance, self.potential)
         expected_output = self.conductance * (self.E - self.potential)
         np.testing.assert_array_almost_equal(output, expected_output)
 
     def test_CUBA(self):
-        model = bst.nn.CUBA()
+        model = brainstate.nn.CUBA()
         output = model.update(self.conductance)
         expected_output = self.conductance * model.scale
         self.assertTrue(u.math.allclose(output, expected_output))
 
     def test_MgBlock(self):
-        model = bst.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
+        model = brainstate.nn.MgBlock(E=self.E, cc_Mg=self.cc_Mg, alpha=self.alpha, beta=self.beta, V_offset=self.V_offset)
         output = model.update(self.conductance, self.potential)
         norm = (1 + self.cc_Mg / self.beta * jnp.exp(self.alpha * (self.V_offset - self.potential)))
         expected_output = self.conductance * (self.E - self.potential) / norm