brainstate 0.1.1__py2.py3-none-any.whl → 0.1.3__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/mixin.py

@@ -132,6 +132,7 @@ class AlignPost(Mixin):
         raise NotImplementedError
 
 
+
 class BindCondData(Mixin):
     """Bind temporary conductance data.
 
@@ -147,7 +148,6 @@ class BindCondData(Mixin):
 
 
 class UpdateReturn(Mixin):
-
     def update_return(self) -> PyTree:
         """
         The update function return of the model.
@@ -157,19 +157,6 @@ class UpdateReturn(Mixin):
         """
         raise NotImplementedError(f'Must implement the "{self.update_return.__name__}()" function.')
 
-    def update_return_info(self) -> PyTree:
-        """
-        The update return information of the model.
-
-        It should be a pytree, with each element as a ``jax.Array``.
-
-        .. note::
-           Should not include the batch axis and batch in_size.
-           These information will be inferred from the ``mode`` attribute.
-
-        """
-        raise NotImplementedError(f'Must implement the "{self.update_return_info.__name__}()" function.')
-
 
 class _MetaUnionType(type):
     def __new__(cls, name, bases, dct):

brainstate/nn/__init__.py

@@ -19,46 +19,110 @@ from ._collective_ops import *
 from ._collective_ops import __all__ as collective_ops_all
 from ._common import *
 from ._common import __all__ as common_all
-from ._dyn_impl import *
-from ._dyn_impl import __all__ as dyn_impl_all
+from ._conv import *
+from ._conv import __all__ as conv_all
+from ._delay import *
+from ._delay import __all__ as state_delay_all
+from ._dropout import *
+from ._dropout import __all__ as dropout_all
 from ._dynamics import *
-from ._dynamics import __all__ as dynamics_all
+from ._dynamics import __all__ as dyn_all
 from ._elementwise import *
 from ._elementwise import __all__ as elementwise_all
-from ._event import *
-from ._event import __all__ as _event_all
+from ._embedding import *
+from ._embedding import __all__ as embed_all
 from ._exp_euler import *
 from ._exp_euler import __all__ as exp_euler_all
-from ._interaction import *
-from ._interaction import __all__ as interaction_all
+from ._fixedprob_mv import EventFixedProb, EventFixedNumConn
+from ._inputs import *
+from ._inputs import __all__ as inputs_all
+from ._linear import *
+from ._linear import __all__ as linear_all
+from ._linear_mv import EventLinear
+from ._ltp import *
+from ._ltp import __all__ as ltp_all
 from ._module import *
 from ._module import __all__ as module_all
+from ._neuron import *
+from ._neuron import __all__ as dyn_neuron_all
+from ._normalizations import *
+from ._normalizations import __all__ as normalizations_all
+from ._poolings import *
+from ._poolings import __all__ as poolings_all
+from ._projection import *
+from ._projection import __all__ as projection_all
+from ._rate_rnns import *
+from ._rate_rnns import __all__ as rate_rnns
+from ._readout import *
+from ._readout import __all__ as readout_all
+from ._stp import *
+from ._stp import __all__ as stp_all
+from ._synapse import *
+from ._synapse import __all__ as dyn_synapse_all
+from ._synaptic_projection import *
+from ._synaptic_projection import __all__ as _syn_proj_all
+from ._synouts import *
+from ._synouts import __all__ as synouts_all
 from ._utils import *
 from ._utils import __all__ as utils_all
 
 __all__ = (
-    ['metrics']
+    [
+        'metrics',
+        'EventLinear',
+        'EventFixedProb',
+        'EventFixedNumConn',
+    ]
     + collective_ops_all
     + common_all
-    + dyn_impl_all
-    + dynamics_all
     + elementwise_all
     + module_all
     + exp_euler_all
-    + interaction_all
     + utils_all
-    + _event_all
+    + dyn_all
+    + projection_all
+    + state_delay_all
+    + synouts_all
+    + conv_all
+    + linear_all
+    + normalizations_all
+    + poolings_all
+    + embed_all
+    + dropout_all
+    + elementwise_all
+    + dyn_neuron_all
+    + dyn_synapse_all
+    + inputs_all
+    + rate_rnns
+    + readout_all
+    + stp_all
+    + ltp_all
+    + _syn_proj_all
 )
 
 del (
     collective_ops_all,
     common_all,
-    dyn_impl_all,
-    dynamics_all,
-    elementwise_all,
     module_all,
     exp_euler_all,
-    interaction_all,
     utils_all,
-    _event_all,
+    dyn_all,
+    projection_all,
+    state_delay_all,
+    synouts_all,
+    conv_all,
+    linear_all,
+    normalizations_all,
+    poolings_all,
+    embed_all,
+    dropout_all,
+    elementwise_all,
+    dyn_neuron_all,
+    dyn_synapse_all,
+    inputs_all,
+    readout_all,
+    rate_rnns,
+    stp_all,
+    ltp_all,
+    _syn_proj_all,
 )

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/{_interaction/_conv.py → _conv.py}

@@ -23,8 +23,8 @@ import jax.numpy as jnp
 
 from brainstate import init, functional
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._module import Module
 
 T = TypeVar('T')
 

brainstate/nn/{_interaction/_conv_test.py → _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/{_dynamics/_state_delay.py → _delay.py}

@@ -27,9 +27,9 @@ from brainstate import environ
 from brainstate._state import ShortTermState, State
 from brainstate.compile import jit_error_if
 from brainstate.graph import Node
-from brainstate.nn._collective_ops import call_order
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, PyTree
+from ._collective_ops import call_order
+from ._module import Module
 
 __all__ = [
     'Delay', 'DelayAccess', 'StateWithDelay',
@@ -135,6 +135,7 @@ class Delay(Module):
         entries: Optional[Dict] = None,  # delay access entry
         delay_method: Optional[str] = _DELAY_ROTATE,  # delay method
         interp_method: str = _INTERP_LINEAR,  # interpolation method
+        take_aware_unit: bool = False
     ):
         # target information
         self.target_info = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, a.dtype), target_info)
@@ -170,6 +171,9 @@ class Delay(Module):
             for entry, delay_time in entries.items():
                 self.register_entry(entry, delay_time)
 
+        self.take_aware_unit = take_aware_unit
+        self._unit = None
+
     @property
     def history(self):
         return self._history

brainstate/nn/{_elementwise/_dropout.py → _dropout.py}

@@ -22,8 +22,8 @@ import jax.numpy as jnp
 
 from brainstate import random, environ, init
 from brainstate._state import ShortTermState
-from brainstate.nn._module import ElementWiseBlock
 from brainstate.typing import Size
+from ._module import ElementWiseBlock
 
 __all__ = [
     'DropoutFixed', 'Dropout', 'Dropout1d', 'Dropout2d', 'Dropout3d',