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

brainstate/nn/{_elementwise/_dropout_test.py → _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/{_dynamics/_dynamics_base.py → _dynamics.py}

@@ -41,13 +41,14 @@ import numpy as np
 from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
-from brainstate.mixin import ParamDescriber
-from brainstate.nn._module import Module
+from brainstate.mixin import ParamDescriber, UpdateReturn
 from brainstate.typing import Size, ArrayLike
-from ._state_delay import StateWithDelay, Delay
+from ._delay import StateWithDelay, Delay
+from ._module import Module
 
 __all__ = [
-    'DynamicsGroup', 'Projection', 'Dynamics', 'Prefetch',
+    'DynamicsGroup', 'Projection', 'Dynamics',
+    'Prefetch', 'PrefetchDelay', 'PrefetchDelayAt', 'OutputDelayAt',
 ]
 
 T = TypeVar('T')
@@ -99,7 +100,7 @@ class Projection(Module):
             raise ValueError('Do not implement the update() function.')
 
 
-class Dynamics(Module):
+class Dynamics(Module, UpdateReturn):
     """
     Base class for implementing neural dynamics models in BrainState.
 
@@ -821,6 +822,41 @@ class Dynamics(Module):
         else:
             raise TypeError(f'The input {dyn} should be an instance of {Dynamics} or a delayed initializer.')
 
+    def prefetch_delay(self, state: str, delay: Optional[ArrayLike] = None) -> 'PrefetchDelayAt':
+        """
+        Create a reference to a delayed state or variable in the module.
+
+        This method simplifies the process of accessing a delayed version of a state or variable
+        within the module. It first creates a prefetch reference to the specified state,
+        then specifies the delay time for accessing this state.
+
+        Args:
+            state (str): The name of the state or variable to reference.
+            delay (Optional[ArrayLike]): The amount of time to delay the variable access,
+                typically in time units (e.g., milliseconds). Defaults to None.
+
+        Returns:
+            PrefetchDelayAt: An object that provides access to the variable at the specified delay time.
+        """
+        return self.prefetch(state).delay.at(delay)
+
+    def output_delay(self, delay: Optional[ArrayLike] = None) -> 'OutputDelayAt':
+        """
+        Create a reference to the delayed output of the module.
+
+        This method simplifies the process of accessing a delayed version of the module's output.
+        It instantiates an `OutputDelayAt` object, which can be used to retrieve the output value
+        at the specified delay time.
+
+        Args:
+            delay (Optional[ArrayLike]): The amount of time to delay the output access,
+                typically in time units (e.g., milliseconds). Defaults to None.
+
+        Returns:
+            OutputDelayAt: An object that provides access to the module's output at the specified delay time.
+        """
+        return OutputDelayAt(self, delay)
+
 
 class Prefetch(Node):
     """
@@ -885,6 +921,7 @@ class Prefetch(Node):
             An object that provides access to delayed versions of the prefetched item.
         """
         return PrefetchDelay(self.module, self.item)
+        # return PrefetchDelayAt(self.module, self.item, time)
 
     def __call__(self, *args, **kwargs):
         """
@@ -1007,7 +1044,7 @@ class PrefetchDelayAt(Node):
         self,
         module: Dynamics,
         item: str,
-        time: ArrayLike
+        time: ArrayLike = None,
     ):
         """
         Initialize a PrefetchDelayAt object.
@@ -1026,14 +1063,16 @@ class PrefetchDelayAt(Node):
         self.module = module
         self.item = item
         self.time = time
-        self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
 
-        # register the delay
-        key = _get_delay_key(item)
-        if not module._has_after_update(key):
-            module._add_after_update(key, not_receive_update_output(StateWithDelay(module, item)))
-        self.state_delay: StateWithDelay = module._get_after_update(key)
-        self.state_delay.register_delay(time)
+        if time is not None:
+            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
+
+            # register the delay
+            key = _get_prefetch_delay_key(item)
+            if not module._has_after_update(key):
+                module._add_after_update(key, not_receive_update_output(StateWithDelay(module, item)))
+            self.state_delay: StateWithDelay = module._get_after_update(key)
+            self.state_delay.register_delay(time)
 
     def __call__(self, *args, **kwargs):
         """
@@ -1044,12 +1083,94 @@ class PrefetchDelayAt(Node):
         Any
             The value of the state or variable at the specified delay time.
         """
-        # return self.state_delay.retrieve_at_time(self.time)
-        return self.state_delay.retrieve_at_step(self.step)
+        if self.time is None:
+            return _get_prefetch_item(self).value
+        else:
+            return self.state_delay.retrieve_at_step(self.step)
+
+
+class OutputDelayAt(Node):
+    """
+    Provides access to a specific delayed state or variable value at the specific time.
+
+    This class represents the final step in the prefetch delay chain, providing
+    actual access to state values at a specific delay time. It converts the
+    specified time delay into steps and registers the delay with the appropriate
+    StateWithDelay handler.
+
+    Parameters
+    ----------
+    module : Dynamics
+        The dynamics module that contains the referenced state or variable.
+    item : str
+        The name of the state or variable to access with delay.
+    time : ArrayLike
+        The amount of time to delay access by, typically in time units (e.g., milliseconds).
+
+    Examples
+    --------
+    >>> import brainstate
+    >>> import brainunit as u
+    >>> neuron = brainstate.nn.LIF(10)
+    >>> # Create a reference to voltage delayed by 5ms
+    >>> delayed_v = PrefetchDelayAt(neuron, 'V', 5.0 * u.ms)
+    >>> # Get the delayed value
+    >>> v_value = delayed_v()
+    """
+
+    def __init__(
+        self,
+        module: Dynamics,
+        time: Optional[ArrayLike] = None,
+    ):
+        """
+        Initialize a PrefetchDelayAt object.
+
+        Parameters
+        ----------
+        module : AlignPre, Module
+            The dynamics module that contains the referenced state or variable.
+        time : ArrayLike
+            The amount of time to delay access by, typically in time units.
+        """
+        super().__init__()
+        assert isinstance(module, UpdateReturn), 'The module should implement the `update_return` method.'
+        assert isinstance(module, Module), 'The module should be an instance of Module.'
+        self.module = module
+        self.time = time
+        if time is not None:
+            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
+
+            # register the delay
+            key = _get_output_delay_key()
+            if not module._has_after_update(key):
+                # TODO: unit processing
+                delay = Delay(module.update_return(), time)
+                module._add_after_update(key, receive_update_output(delay))
+            self.out_delay: Delay = module._get_after_update(key)
+            self.out_delay.register_delay(time)
+
+    def __call__(self, *args, **kwargs):
+        """
+        Retrieve the value of the state at the specified delay time.
+
+        Returns
+        -------
+        Any
+            The value of the state or variable at the specified delay time.
+        """
+        if self.time is None:
+            return self.module.update_return()
+        else:
+            return self.out_delay.retrieve_at_step(self.step)
+
+
+def _get_prefetch_delay_key(item) -> str:
+    return f'{item}-prefetch-delay'
 
 
-def _get_delay_key(item) -> str:
-    return f'{item}-delay'
+def _get_output_delay_key() -> str:
+    return f'output-delay'
 
 
 def _get_prefetch_item(target: Union[Prefetch, PrefetchDelayAt]) -> Any:
@@ -1064,7 +1185,7 @@ def _get_prefetch_item_delay(target: Union[Prefetch, PrefetchDelay, PrefetchDela
         f'The target module should be an instance '
         f'of Dynamics. But got {target.module}.'
     )
-    delay = target.module._get_after_update(_get_delay_key(target.item))
+    delay = target.module._get_after_update(_get_prefetch_delay_key(target.item))
     if not isinstance(delay, StateWithDelay):
         raise TypeError(f'The prefetch target should be a {StateWithDelay.__name__} when accessing '
                         f'its delay. But got {delay}.')

brainstate/nn/{_dynamics/_dynamics_base_test.py → _dynamics_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/{_elementwise/_elementwise.py → _elementwise.py}

@@ -23,8 +23,8 @@ import jax.typing
 
 from brainstate import random, functional as F
 from brainstate._state import ParamState
-from brainstate.nn._module import ElementWiseBlock
 from brainstate.typing import ArrayLike
+from ._module import ElementWiseBlock
 
 __all__ = [
     # activation functions

brainstate/nn/_elementwise_test.py

@@ -0,0 +1,169 @@
+# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+from absl.testing import absltest
+from absl.testing import parameterized
+
+import brainstate
+
+
+class Test_Activation(parameterized.TestCase):
+
+    def test_Threshold(self):
+        threshold_layer = brainstate.nn.Threshold(5, 20)
+        input = brainstate.random.randn(2)
+        output = threshold_layer(input)
+
+    def test_ReLU(self):
+        ReLU_layer = brainstate.nn.ReLU()
+        input = brainstate.random.randn(2)
+        output = ReLU_layer(input)
+
+    def test_RReLU(self):
+        RReLU_layer = brainstate.nn.RReLU(lower=0, upper=1)
+        input = brainstate.random.randn(2)
+        output = RReLU_layer(input)
+
+    def test_Hardtanh(self):
+        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 = brainstate.nn.ReLU6()
+        input = brainstate.random.randn(2)
+        output = ReLU6_layer(input)
+
+    def test_Sigmoid(self):
+        Sigmoid_layer = brainstate.nn.Sigmoid()
+        input = brainstate.random.randn(2)
+        output = Sigmoid_layer(input)
+
+    def test_Hardsigmoid(self):
+        Hardsigmoid_layer = brainstate.nn.Hardsigmoid()
+        input = brainstate.random.randn(2)
+        output = Hardsigmoid_layer(input)
+
+    def test_Tanh(self):
+        Tanh_layer = brainstate.nn.Tanh()
+        input = brainstate.random.randn(2)
+        output = Tanh_layer(input)
+
+    def test_SiLU(self):
+        SiLU_layer = brainstate.nn.SiLU()
+        input = brainstate.random.randn(2)
+        output = SiLU_layer(input)
+
+    def test_Mish(self):
+        Mish_layer = brainstate.nn.Mish()
+        input = brainstate.random.randn(2)
+        output = Mish_layer(input)
+
+    def test_Hardswish(self):
+        Hardswish_layer = brainstate.nn.Hardswish()
+        input = brainstate.random.randn(2)
+        output = Hardswish_layer(input)
+
+    def test_ELU(self):
+        ELU_layer = brainstate.nn.ELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
+        output = ELU_layer(input)
+
+    def test_CELU(self):
+        CELU_layer = brainstate.nn.CELU(alpha=0.5, )
+        input = brainstate.random.randn(2)
+        output = CELU_layer(input)
+
+    def test_SELU(self):
+        SELU_layer = brainstate.nn.SELU()
+        input = brainstate.random.randn(2)
+        output = SELU_layer(input)
+
+    def test_GLU(self):
+        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 = brainstate.nn.GELU()
+        input = brainstate.random.randn(2)
+        output = GELU_layer(input)
+
+    def test_Hardshrink(self):
+        Hardshrink_layer = brainstate.nn.Hardshrink(lambd=1)
+        input = brainstate.random.randn(2)
+        output = Hardshrink_layer(input)
+
+    def test_LeakyReLU(self):
+        LeakyReLU_layer = brainstate.nn.LeakyReLU()
+        input = brainstate.random.randn(2)
+        output = LeakyReLU_layer(input)
+
+    def test_LogSigmoid(self):
+        LogSigmoid_layer = brainstate.nn.LogSigmoid()
+        input = brainstate.random.randn(2)
+        output = LogSigmoid_layer(input)
+
+    def test_Softplus(self):
+        Softplus_layer = brainstate.nn.Softplus()
+        input = brainstate.random.randn(2)
+        output = Softplus_layer(input)
+
+    def test_Softshrink(self):
+        Softshrink_layer = brainstate.nn.Softshrink(lambd=1)
+        input = brainstate.random.randn(2)
+        output = Softshrink_layer(input)
+
+    def test_PReLU(self):
+        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 = brainstate.nn.Softsign()
+        input = brainstate.random.randn(2)
+        output = Softsign_layer(input)
+
+    def test_Tanhshrink(self):
+        Tanhshrink_layer = brainstate.nn.Tanhshrink()
+        input = brainstate.random.randn(2)
+        output = Tanhshrink_layer(input)
+
+    def test_Softmin(self):
+        Softmin_layer = brainstate.nn.Softmin(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = Softmin_layer(input)
+
+    def test_Softmax(self):
+        Softmax_layer = brainstate.nn.Softmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = Softmax_layer(input)
+
+    def test_Softmax2d(self):
+        Softmax2d_layer = brainstate.nn.Softmax2d()
+        input = brainstate.random.randn(2, 3, 12, 13)
+        output = Softmax2d_layer(input)
+
+    def test_LogSoftmax(self):
+        LogSoftmax_layer = brainstate.nn.LogSoftmax(dim=2)
+        input = brainstate.random.randn(2, 3, 4)
+        output = LogSoftmax_layer(input)
+
+
+if __name__ == '__main__':
+    absltest.main()

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

@@ -17,8 +17,8 @@ from typing import Optional, Callable, Union
 
 from brainstate import init
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike
+from ._module import Module
 
 __all__ = [
     'Embedding',

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/{_event/_fixedprob_mv.py → _fixedprob_mv.py}

@@ -25,8 +25,8 @@ from brainstate import random, augment, environ, init
 from brainstate._compatible_import import brainevent
 from brainstate._state import ParamState
 from brainstate.compile import for_loop
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
     'EventFixedNumConn',

brainstate/nn/{_event/_fixedprob_mv_test.py → _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/{_dyn_impl/_inputs.py → _inputs.py}

@@ -20,12 +20,11 @@ import jax
 import numpy as np
 
 from brainstate import environ, init, random
-from brainstate._state import ShortTermState
-from brainstate._state import State, maybe_state
+from brainstate._state import ShortTermState, State, maybe_state
 from brainstate.compile import while_loop
-from brainstate.nn._dynamics._dynamics_base import Dynamics, Prefetch
-from brainstate.nn._module import Module
 from brainstate.typing import ArrayLike, Size, DTypeLike
+from ._dynamics import Dynamics, Prefetch
+from ._module import Module
 
 __all__ = [
     'SpikeTime',
@@ -134,7 +133,7 @@ class PoissonSpike(Dynamics):
         self.freqs = init.param(freqs, self.varshape, allow_none=False)
 
     def update(self):
-        spikes = random.rand(self.varshape) <= (self.freqs * environ.get_dt())
+        spikes = random.rand(*self.varshape) <= (self.freqs * environ.get_dt())
         spikes = u.math.asarray(spikes, dtype=self.spk_type)
         return spikes
 

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

@@ -22,8 +22,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, Size
+from ._module import Module
 
 __all__ = [
     'Linear',
@@ -350,10 +350,7 @@ class OneToOne(Module):
         self.weight = param_type(param)
 
     def update(self, pre_val):
-        pre_val, pre_unit = u.get_mantissa(pre_val), u.get_unit(pre_val)
-        w_val, w_unit = u.get_mantissa(self.weight.value['weight']), u.get_unit(self.weight.value['weight'])
-        post_val = pre_val * w_val
-        post_val = u.maybe_decimal(u.Quantity(post_val, unit=w_unit * pre_unit))
+        post_val = pre_val * self.weight.value['weight']
         if 'bias' in self.weight.value:
             post_val = post_val + self.weight.value['bias']
         return post_val

brainstate/nn/{_event/_linear_mv.py → _linear_mv.py}

@@ -21,8 +21,8 @@ import jax
 from brainstate import init
 from brainstate._compatible_import import brainevent
 from brainstate._state import ParamState
-from brainstate.nn._module import Module
 from brainstate.typing import Size, ArrayLike
+from ._module import Module
 
 __all__ = [
     'EventLinear',

brainstate/nn/{_event/_linear_mv_test.py → _linear_mv_test.py}

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