brainstate 0.1.7__py2.py3-none-any.whl → 0.1.8__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.7"
+__version__ = "0.1.8"
 
 from . import augment
 from . import compile

brainstate/compile/_make_jaxpr.py

@@ -352,11 +352,13 @@ class StatefulFunction(PrettyObject):
             cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_write_states()
 
-    def _check_input(self, x):
+    def _check_input_ouput(self, x):
         if isinstance(x, State):
-            raise ValueError(
-                'Inputs for brainstate transformations cannot be an instance of State. '
-                f'But we got {x}'
+            x.raise_error_with_source_info(
+                ValueError(
+                    'Inputs/outputs for brainstate transformations cannot be an instance of State. '
+                    f'But we got {x}'
+                )
             )
 
     def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
@@ -477,9 +479,7 @@ class StatefulFunction(PrettyObject):
 
         # State instance as functional returns is not allowed.
         # Checking whether the states are returned.
-        for leaf in jax.tree.leaves(out):
-            if isinstance(leaf, State):
-                leaf.raise_error_with_source_info(ValueError(f"State object is not allowed to be returned: {leaf}"))
+        jax.tree.map(self._check_input_ouput, out, is_leaf=lambda x: isinstance(x, State))
         return out, state_values
 
     def make_jaxpr(self, *args, return_only_write: bool = False, **kwargs):
@@ -501,7 +501,7 @@ class StatefulFunction(PrettyObject):
         cache_key = self.get_arg_cache_key(*args, **kwargs)
 
         # check input types
-        jax.tree.map(self._check_input, (args, kwargs), is_leaf=lambda x: isinstance(x, State))
+        jax.tree.map(self._check_input_ouput, (args, kwargs), is_leaf=lambda x: isinstance(x, State))
 
         if cache_key not in self._cached_state_trace:
             try:

brainstate/compile/_make_jaxpr_test.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-from __future__ import annotations
 
 import unittest
 
@@ -132,15 +131,26 @@ class TestMakeJaxpr(unittest.TestCase):
         print(jaxpr)
         print(states)
 
+    def test_state_in(self):
+        def f(a):
+            return a.value
 
-def test_return_states():
-    import jax.numpy
+        with pytest.raises(ValueError):
+            brainstate.compile.StatefulFunction(f).make_jaxpr(brainstate.State(1.))
 
-    a = brainstate.State(jax.numpy.ones(3))
+    def test_state_out(self):
+        def f(a):
+            return brainstate.State(a)
 
-    @brainstate.compile.jit
-    def f():
-        return a
+        with pytest.raises(ValueError):
+            brainstate.compile.StatefulFunction(f).make_jaxpr(1.)
 
-    with pytest.raises(ValueError):
-        f()
+    def test_return_states(self):
+        a = brainstate.State(jnp.ones(3))
+
+        @brainstate.compile.jit
+        def f():
+            return a
+
+        with pytest.raises(ValueError):
+            f()

brainstate/nn/_delay.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 # ==============================================================================
 
-import math
 import numbers
 from functools import partial
 from typing import Optional, Dict, Callable, Union, Sequence
@@ -41,34 +40,27 @@ _INTERP_LINEAR = 'linear_interp'
 _INTERP_ROUND = 'round'
 
 
-def _get_delay(delay_time, delay_step):
-    if delay_time is None:
-        if delay_step is None:
-            return 0., 0
-        else:
-            assert isinstance(delay_step, int), '"delay_step" should be an integer.'
-            if delay_step == 0:
-                return 0., 0
-            with jax.ensure_compile_time_eval():
-                delay_time = delay_step * environ.get_dt()
-    else:
-        assert delay_step is None, '"delay_step" should be None if "delay_time" is given.'
-        # assert isinstance(delay_time, (int, float))
-        with jax.ensure_compile_time_eval():
-            delay_step = delay_time / environ.get_dt()
-        delay_step = math.ceil(float(delay_step))
+def _get_delay(delay_time):
+    with jax.ensure_compile_time_eval():
+        if delay_time is None:
+            return 0. * environ.get_dt(), 0
+        delay_step = delay_time / environ.get_dt()
+        delay_step = jnp.ceil(delay_step).astype(environ.ditype())
     return delay_time, delay_step
 
 
 class DelayAccess(Node):
     """
-    The delay access class.
+    Accessor node for a registered entry in a Delay instance.
+
+    This node holds a reference to a Delay and a named entry that was
+    registered on that Delay. It is used by graphs to query delayed
+    values by delegating to the underlying Delay instance.
 
     Args:
-      delay: The delay instance.
-      time: The delay time.
-      indices: The indices of the delay data.
-      delay_entry: The delay entry.
+        delay: The delay instance.
+        time: The delay time.
+        delay_entry: The delay entry.
     """
 
     __module__ = 'brainstate.nn'
@@ -78,17 +70,15 @@ class DelayAccess(Node):
         delay: 'Delay',
         time: Union[None, int, float],
         delay_entry: str,
-        *indices,
     ):
         super().__init__()
         self.refs = {'delay': delay}
         assert isinstance(delay, Delay), 'The input delay should be an instance of Delay.'
         self._delay_entry = delay_entry
-        delay.register_entry(self._delay_entry, time)
-        self.indices = indices
+        self.delay_info = delay.register_entry(self._delay_entry, time)
 
     def update(self):
-        return self.refs['delay'].at(self._delay_entry, *self.indices)
+        return self.refs['delay'].at(self._delay_entry)
 
 
 class Delay(Module):
@@ -141,17 +131,21 @@ class Delay(Module):
         self.target_info = jax.tree.map(lambda a: jax.ShapeDtypeStruct(a.shape, a.dtype), target_info)
 
         # delay method
-        assert delay_method in [_DELAY_ROTATE, _DELAY_CONCAT], (f'Un-supported delay method {delay_method}. '
-                                                                f'Only support {_DELAY_ROTATE} and {_DELAY_CONCAT}')
+        assert delay_method in [_DELAY_ROTATE, _DELAY_CONCAT], (
+            f'Un-supported delay method {delay_method}. '
+            f'Only support {_DELAY_ROTATE} and {_DELAY_CONCAT}'
+        )
         self.delay_method = delay_method
 
         # interp method
-        assert interp_method in [_INTERP_LINEAR, _INTERP_ROUND], (f'Un-supported interpolation method {interp_method}. '
-                                                                  f'we only support: {[_INTERP_LINEAR, _INTERP_ROUND]}')
+        assert interp_method in [_INTERP_LINEAR, _INTERP_ROUND], (
+            f'Un-supported interpolation method {interp_method}. '
+            f'we only support: {[_INTERP_LINEAR, _INTERP_ROUND]}'
+        )
         self.interp_method = interp_method
 
         # delay length and time
-        self.max_time, delay_length = _get_delay(time, None)
+        self.max_time, delay_length = _get_delay(time)
         self.max_length = delay_length + 1
 
         super().__init__()
@@ -169,7 +163,10 @@ class Delay(Module):
         # other info
         if entries is not None:
             for entry, delay_time in entries.items():
-                self.register_entry(entry, delay_time)
+                if isinstance(delay_time, (tuple, list)):
+                    self.register_entry(entry, *delay_time)
+                else:
+                    self.register_entry(entry, delay_time)
 
         self.take_aware_unit = take_aware_unit
         self._unit = None
@@ -205,73 +202,107 @@ class Delay(Module):
         fun = partial(self._f_to_init, length=self.max_length, batch_size=batch_size)
         self.history.value = jax.tree.map(fun, self.target_info)
 
-    def register_delay(
-        self,
-        delay_time: Optional[Union[int, float]] = None,
-        delay_step: Optional[int] = None,
-    ):
-        if isinstance(delay_time, (np.ndarray, jax.Array)):
-            assert delay_time.size == 1 and delay_time.ndim == 0
-            delay_time = delay_time.item()
+    def register_delay(self, *delay_time):
+        """
+        Register delay times and update the maximum delay configuration.
+
+        This method processes one or more delay times, validates their format and consistency,
+        and updates the delay buffer size if necessary. It handles both scalar and vector
+        delay times, ensuring all vector delays have the same size.
 
-        _, delay_step = _get_delay(delay_time, delay_step)
+        Args:
+            *delay_time: Variable number of delay time arguments. The first argument should be
+                the primary delay time (float, int, or array-like). Additional arguments are
+                treated as indices or secondary delay parameters. All delay times should be
+                non-negative numbers or arrays of the same size.
+
+        Returns:
+            tuple or None: If delay_time[0] is None, returns None. Otherwise, returns a tuple
+                containing (delay_step, *delay_time[1:]) where delay_step is the computed
+                delay step in integer time units, and the remaining elements are the
+                additional delay parameters passed in.
+
+        Raises:
+            AssertionError: If no delay time is provided (empty delay_time).
+            ValueError: If delay times have inconsistent sizes when using vector delays,
+                or if delay times are not scalar or 1D arrays.
+
+        Note:
+            - The method updates self.max_time and self.max_length if the new delay
+              requires a larger buffer size.
+            - Delay steps are computed using the current environment time step (dt).
+            - All delay indices (delay_time[1:]) must be integers.
+            - Vector delays must all have the same size as the first delay time.
+
+        Example:
+            >>> delay_obj.register_delay(5.0)  # Register 5ms delay
+            >>> delay_obj.register_delay(jnp.array([2.0, 3.0]), 0, 1)  # Vector delay with indices
+        """
+        assert len(delay_time) >= 1, 'You should provide at least one delay time.'
+        delay_size = u.math.size(delay_time[0])
+        for dt in delay_time[1:]:
+            assert jnp.issubdtype(u.math.get_dtype(dt), jnp.integer), f'The index should be integer. But got {dt}.'
+        for dt in delay_time:
+            if u.math.ndim(dt) == 0:
+                pass
+            elif u.math.ndim(dt) == 1:
+                if u.math.size(dt) != delay_size:
+                    raise ValueError(
+                        f'The delay time should be a scalar or a vector with the same size. '
+                        f'But got {delay_time}. The delay time {dt} has size {u.math.size(dt)}'
+                    )
+            else:
+                raise ValueError(f'The delay time should be a scalar/vector. But got {dt}.')
+        if delay_time[0] is None:
+            return None
+        time, delay_step = _get_delay(delay_time[0])
+        max_delay_step = jnp.max(delay_step)
+        self.max_time = u.math.max(time)
 
         # delay variable
-        if self.max_length <= delay_step + 1:
-            self.max_length = delay_step + 1
-            self.max_time = delay_time
-        return self
+        if self.max_length <= max_delay_step + 1:
+            self.max_length = max_delay_step + 1
+        return delay_step, *delay_time[1:]
 
-    def register_entry(
-        self,
-        entry: str,
-        delay_time: Optional[Union[int, float]] = None,
-        delay_step: Optional[int] = None,
-    ) -> 'Delay':
+    def register_entry(self, entry: str, *delay_time) -> 'Delay':
         """
         Register an entry to access the delay data.
 
         Args:
-          entry: str. The entry to access the delay data.
-          delay_time: The delay time of the entry (can be a float).
-          delay_step: The delay step of the entry (must be an int). ``delat_step = delay_time / dt``.
-
-        Returns:
-          Return the self.
+            entry: str. The entry to access the delay data.
+            delay_time: The delay time of the entry, the first element is the delay time,
+                the second and later element is the index.
         """
         if entry in self._registered_entries:
-            raise KeyError(f'Entry {entry} has been registered. '
-                           f'The existing delay for the key {entry} is {self._registered_entries[entry]}. '
-                           f'The new delay for the key {entry} is {delay_time}. '
-                           f'You can use another key. ')
-
-        if isinstance(delay_time, (np.ndarray, jax.Array)):
-            assert delay_time.size == 1 and delay_time.ndim == 0
-            delay_time = delay_time.item()
+            raise KeyError(
+                f'Entry {entry} has been registered. '
+                f'The existing delay for the key {entry} is {self._registered_entries[entry]}. '
+                f'The new delay for the key {entry} is {delay_time}. '
+                f'You can use another key. '
+            )
+        delay_info = self.register_delay(*delay_time)
+        self._registered_entries[entry] = delay_info
+        return delay_info
 
-        _, delay_step = _get_delay(delay_time, delay_step)
+    def access(self, entry: str, delay_time: Sequence) -> DelayAccess:
+        """
+        Create a DelayAccess object for a specific delay entry and delay time.
 
-        # delay variable
-        if self.max_length <= delay_step + 1:
-            self.max_length = delay_step + 1
-            self.max_time = delay_time
-        self._registered_entries[entry] = delay_step
-        return self
+        Args:
+            entry (str): The name of the delay entry to access.
+            delay_time (Sequence): The delay time or parameters associated with the entry.
 
-    def access(
-        self,
-        entry: str = None,
-        time: Sequence = None,
-    ) -> DelayAccess:
-        return DelayAccess(self, time, delay_entry=entry)
+        Returns:
+            DelayAccess: An object that provides access to the delay data for the specified entry and time.
+        """
+        return DelayAccess(self, delay_time, delay_entry=entry)
 
-    def at(self, entry: str, *indices) -> ArrayLike:
+    def at(self, entry: str) -> ArrayLike:
         """
         Get the data at the given entry.
 
         Args:
           entry: str. The entry to access the data.
-          *indices: The slicing indices. Not include the slice at the batch dimension.
 
         Returns:
           The data.
@@ -282,8 +313,8 @@ class Delay(Module):
             raise KeyError(f'Does not find delay entry "{entry}".')
         delay_step = self._registered_entries[entry]
         if delay_step is None:
-            delay_step = 0
-        return self.retrieve_at_step(delay_step, *indices)
+            delay_step = (0,)
+        return self.retrieve_at_step(*delay_step)
 
     def retrieve_at_step(self, delay_step, *indices) -> PyTree:
         """
@@ -306,8 +337,10 @@ class Delay(Module):
 
         if environ.get(environ.JIT_ERROR_CHECK, False):
             def _check_delay(delay_len):
-                raise ValueError(f'The request delay length should be less than the '
-                                 f'maximum delay {self.max_length - 1}. But we got {delay_len}')
+                raise ValueError(
+                    f'The request delay length should be less than the '
+                    f'maximum delay {self.max_length - 1}. But we got {delay_len}'
+                )
 
             jit_error_if(delay_step >= self.max_length, _check_delay, delay_step)
 
@@ -363,13 +396,17 @@ class Delay(Module):
 
         if environ.get(environ.JIT_ERROR_CHECK, False):
             def _check_delay(t_now, t_delay):
-                raise ValueError(f'The request delay time should be within '
-                                 f'[{t_now - self.max_time - dt}, {t_now}], '
-                                 f'but we got {t_delay}')
+                raise ValueError(
+                    f'The request delay time should be within '
+                    f'[{t_now - self.max_time - dt}, {t_now}], '
+                    f'but we got {t_delay}'
+                )
 
             jit_error_if(
-                jnp.logical_or(delay_time > current_time,
-                               delay_time < current_time - self.max_time - dt),
+                jnp.logical_or(
+                    delay_time > current_time,
+                    delay_time < current_time - self.max_time - dt
+                ),
                 _check_delay,
                 current_time,
                 delay_time
@@ -385,10 +422,7 @@ class Delay(Module):
             return jax.tree.map(lambda a, b: a * (1 - t_diff) + b * t_diff, data_at_t0, data_at_t1)
 
         elif self.interp_method == _INTERP_ROUND:  # "round" interpolation
-            return self.retrieve_at_step(
-                jnp.asarray(jnp.round(float_time_step), dtype=jnp.int32),
-                *indices
-            )
+            return self.retrieve_at_step(jnp.asarray(jnp.round(float_time_step), dtype=jnp.int32), *indices)
 
         else:  # raise error
             raise ValueError(f'Un-supported interpolation method {self.interp_method}, '
@@ -429,8 +463,6 @@ class Delay(Module):
             raise ValueError(f'Unknown updating method "{self.delay_method}"')
 
 
-
-
 class StateWithDelay(Delay):
     """
     A ``State`` type that defines the state in a differential equation.
@@ -440,8 +472,8 @@ class StateWithDelay(Delay):
 
     state: State  # state
 
-    def __init__(self, target: Node, item: str):
-        super().__init__(None)
+    def __init__(self, target: Node, item: str, init: Callable = None):
+        super().__init__(None, init=init)
 
         self._target = target
         self._target_term = item

brainstate/nn/_delay_test.py

@@ -0,0 +1,184 @@
+# Copyright 2025 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.
+# ==============================================================================
+
+
+import unittest
+
+import jax.numpy as jnp
+
+import brainstate
+
+brainstate.environ.set(dt=0.1)
+
+
+class TestDelay(unittest.TestCase):
+    def test_delay1(self):
+        a = brainstate.State(brainstate.random.random(10, 20))
+        delay = brainstate.nn.Delay(a.value)
+        delay.register_entry('a', 1.)
+        delay.register_entry('b', 2.)
+        delay.register_entry('c', None)
+
+        delay.init_state()
+        with self.assertRaises(KeyError):
+            delay.register_entry('c', 10.)
+
+    def test_rotation_delay(self):
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
+        t0 = 0.
+        t1, n1 = 1., 10
+        t2, n2 = 2., 20
+
+        rotation_delay.register_entry('a', t0)
+        rotation_delay.register_entry('b', t1)
+        rotation_delay.register_entry('c2', 1.9)
+        rotation_delay.register_entry('c', t2)
+
+        rotation_delay.init_state()
+
+        print()
+        # print(rotation_delay)
+        # print(rotation_delay.max_length)
+
+        for i in range(100):
+            brainstate.environ.set(i=i)
+            rotation_delay.update(jnp.ones((1,)) * i)
+            # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
+            self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
+            self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
+            self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
+
+    def test_concat_delay(self):
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
+        t0 = 0.
+        t1, n1 = 1., 10
+        t2, n2 = 2., 20
+
+        rotation_delay.register_entry('a', t0)
+        rotation_delay.register_entry('b', t1)
+        rotation_delay.register_entry('c', t2)
+
+        rotation_delay.init_state()
+
+        print()
+        for i in range(100):
+            brainstate.environ.set(i=i)
+            rotation_delay.update(jnp.ones((1,)) * i)
+            print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
+            self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
+            self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
+            self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
+        # brainstate.util.clear_buffer_memory()
+
+    def test_jit_erro(self):
+        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
+        rotation_delay.init_state()
+
+        with brainstate.environ.context(i=0, t=0, jit_error_check=True):
+            rotation_delay.retrieve_at_time(-2.0)
+            with self.assertRaises(Exception):
+                rotation_delay.retrieve_at_time(-2.1)
+            rotation_delay.retrieve_at_time(-2.01)
+            with self.assertRaises(Exception):
+                rotation_delay.retrieve_at_time(-2.09)
+            with self.assertRaises(Exception):
+                rotation_delay.retrieve_at_time(0.1)
+            with self.assertRaises(Exception):
+                rotation_delay.retrieve_at_time(0.01)
+
+    def test_round_interp(self):
+        for shape in [(1,), (1, 1), (1, 1, 1)]:
+            for delay_method in ['rotation', 'concat']:
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='round')
+                t0, n1 = 0.01, 0
+                t1, n1 = 1.04, 10
+                t2, n2 = 1.06, 11
+                rotation_delay.init_state()
+
+                @brainstate.compile.jit
+                def retrieve(td, i):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
+                        return rotation_delay.retrieve_at_time(td)
+
+                print()
+                for i in range(100):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
+                        rotation_delay.update(jnp.ones(shape) * i)
+                        print(i,
+                              retrieve(t - t0, i),
+                              retrieve(t - t1, i),
+                              retrieve(t - t2, i))
+                        self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.ones(shape) * i))
+                        self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
+                        self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
+
+    def test_linear_interp(self):
+        for shape in [(1,), (1, 1), (1, 1, 1)]:
+            for delay_method in ['rotation', 'concat']:
+                print(shape, delay_method)
+
+                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
+                                                     interp_method='linear_interp')
+                t0, n0 = 0.01, 0.1
+                t1, n1 = 1.04, 10.4
+                t2, n2 = 1.06, 10.6
+                rotation_delay.init_state()
+
+                @brainstate.compile.jit
+                def retrieve(td, i):
+                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
+                        return rotation_delay.retrieve_at_time(td)
+
+                print()
+                for i in range(100):
+                    t = i * brainstate.environ.get_dt()
+                    with brainstate.environ.context(i=i, t=t):
+                        rotation_delay.update(jnp.ones(shape) * i)
+                        print(i,
+                              retrieve(t - t0, i),
+                              retrieve(t - t1, i),
+                              retrieve(t - t2, i))
+                        self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.maximum(jnp.ones(shape) * i - n0, 0.)))
+                        self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
+                        self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
+
+    def test_rotation_and_concat_delay(self):
+        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
+        concat_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
+        t0 = 0.
+        t1, n1 = 1., 10
+        t2, n2 = 2., 20
+
+        rotation_delay.register_entry('a', t0)
+        rotation_delay.register_entry('b', t1)
+        rotation_delay.register_entry('c', t2)
+        concat_delay.register_entry('a', t0)
+        concat_delay.register_entry('b', t1)
+        concat_delay.register_entry('c', t2)
+
+        rotation_delay.init_state()
+        concat_delay.init_state()
+
+        print()
+        for i in range(100):
+            brainstate.environ.set(i=i)
+            new = jnp.ones((1,)) * i
+            rotation_delay.update(new)
+            concat_delay.update(new)
+            self.assertTrue(jnp.allclose(rotation_delay.at('a'), concat_delay.at('a'), ))
+            self.assertTrue(jnp.allclose(rotation_delay.at('b'), concat_delay.at('b'), ))
+            self.assertTrue(jnp.allclose(rotation_delay.at('c'), concat_delay.at('c'), ))

brainstate/nn/_dynamics.py

@@ -33,9 +33,8 @@ For handling the delays:
 
 """
 
-from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING
+from typing import Any, Dict, Callable, Hashable, Optional, Union, TypeVar, TYPE_CHECKING, Tuple
 
-import brainunit as u
 import jax
 import numpy as np
 
@@ -43,7 +42,7 @@ from brainstate import environ
 from brainstate._state import State
 from brainstate.graph import Node
 from brainstate.mixin import ParamDescriber
-from brainstate.typing import Size, ArrayLike, PyTree
+from brainstate.typing import Size, ArrayLike
 from ._delay import StateWithDelay, Delay
 from ._module import Module
 
@@ -794,10 +793,7 @@ class Dynamics(Module):
         """
         return Prefetch(self, item)
 
-    def align_pre(
-        self,
-        dyn: Union[ParamDescriber[T], T]
-    ) -> T:
+    def align_pre(self, dyn: Union[ParamDescriber[T], T]) -> T:
         """
         Registers a dynamics module to execute after this module.
 
@@ -844,11 +840,7 @@ 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':
+    def prefetch_delay(self, state: str, delay_time, init: Callable = None) -> 'PrefetchDelayAt':
         """
         Create a reference to a delayed state or variable in the module.
 
@@ -858,19 +850,17 @@ class Dynamics(Module):
 
         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.
+            delay_time (ArrayLike): The amount of time to delay the variable access,
+                typically in time units (e.g., milliseconds).
+            init (Callable, optional): An optional initialization function to provide
+                a default value if the delayed state is not yet available.
 
         Returns:
             PrefetchDelayAt: An object that provides access to the variable at the specified delay time.
         """
-        return self.prefetch(state).delay.at(delay)
+        return PrefetchDelayAt(self, state, delay_time, init=init)
 
-    def output_delay(
-        self,
-        delay: Optional[ArrayLike] = None,
-        variable_like: PyTree = None
-    ) -> 'OutputDelayAt':
+    def output_delay(self, *delay_time) -> 'OutputDelayAt':
         """
         Create a reference to the delayed output of the module.
 
@@ -881,12 +871,11 @@ class Dynamics(Module):
         Args:
             delay (Optional[ArrayLike]): The amount of time to delay the output access,
                 typically in time units (e.g., milliseconds). Defaults to None.
-            variable_like:
 
         Returns:
             OutputDelayAt: An object that provides access to the module's output at the specified delay time.
         """
-        return OutputDelayAt(self, delay)
+        return OutputDelayAt(self, delay_time)
 
 
 class Prefetch(Node):
@@ -1024,7 +1013,7 @@ class PrefetchDelay(Node):
         self.module = module
         self.item = item
 
-    def at(self, time: ArrayLike):
+    def at(self, *delay_time):
         """
         Specifies the delay time for accessing the variable.
 
@@ -1039,7 +1028,7 @@ class PrefetchDelay(Node):
         PrefetchDelayAt
             An object that provides access to the variable at the specified delay time.
         """
-        return PrefetchDelayAt(self.module, self.item, time)
+        return PrefetchDelayAt(self.module, self.item, delay_time)
 
 
 class PrefetchDelayAt(Node):
@@ -1075,7 +1064,8 @@ class PrefetchDelayAt(Node):
         self,
         module: Dynamics,
         item: str,
-        time: ArrayLike = None,
+        delay_time: Tuple,
+        init: Callable = None
     ):
         """
         Initialize a PrefetchDelayAt object.
@@ -1086,24 +1076,27 @@ class PrefetchDelayAt(Node):
             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.
+        delay_time : Tuple
+            The amount of time to delay access by, typically in time units (e.g., milliseconds).
         """
         super().__init__()
-        assert isinstance(module, Dynamics), ''
+        assert isinstance(module, Dynamics), 'The module should be an instance of Dynamics.'
         self.module = module
         self.item = item
-        self.time = time
-
-        if time is not None:
-            self.step = u.math.asarray(time / environ.get_dt(), dtype=environ.ditype())
-
-            # register the delay
+        if not isinstance(delay_time, (tuple, list)):
+            delay_time = (delay_time,)
+        self.delay_time = delay_time
+        if len(delay_time) > 0:
             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)))
+                module._add_after_update(
+                    key,
+                    not_receive_update_output(
+                        StateWithDelay(module, item, init=init)
+                    )
+                )
             self.state_delay: StateWithDelay = module._get_after_update(key)
-            self.state_delay.register_delay(time)
+            self.delay_info = self.state_delay.register_delay(*delay_time)
 
     def __call__(self, *args, **kwargs):
         """
@@ -1114,10 +1107,10 @@ class PrefetchDelayAt(Node):
         Any
             The value of the state or variable at the specified delay time.
         """
-        if self.time is None:
+        if len(self.delay_time) == 0:
             return _get_prefetch_item(self).value
         else:
-            return self.state_delay.retrieve_at_step(self.step)
+            return self.state_delay.retrieve_at_step(*self.delay_info)
 
 
 class OutputDelayAt(Node):
@@ -1150,32 +1143,20 @@ class OutputDelayAt(Node):
     def __init__(
         self,
         module: Dynamics,
-        time: Optional[ArrayLike] = None,
-        variable_like: Optional[PyTree] = None,
+        delay_time: Tuple,
     ):
         super().__init__()
         assert isinstance(module, Dynamics), 'The module should be an instance of Dynamics.'
         self.module = module
-        dt = environ.get_dt()
-        if time is None:
-            time = u.math.zeros_like(dt)
-        self.time = time
-        self.step = u.math.asarray(time / dt, dtype=environ.ditype())
-
-        # register the delay
         key = _get_output_delay_key()
         if not module._has_after_update(key):
-            delay = Delay(
-                jax.ShapeDtypeStruct(module.out_size, dtype=environ.dftype()),
-                time,
-                take_aware_unit=True
-            )
+            delay = Delay(jax.ShapeDtypeStruct(module.out_size, dtype=environ.dftype()), take_aware_unit=True)
             module._add_after_update(key, receive_update_output(delay))
         self.out_delay: Delay = module._get_after_update(key)
-        self.out_delay.register_delay(time)
+        self.delay_info = self.out_delay.register_delay(*delay_time)
 
     def __call__(self, *args, **kwargs):
-        return self.out_delay.retrieve_at_step(self.step)
+        return self.out_delay.retrieve_at_step(*self.delay_info)
 
 
 def _get_prefetch_delay_key(item) -> str:
@@ -1242,8 +1223,9 @@ def maybe_init_prefetch(target, *args, **kwargs):
         _get_prefetch_item_delay(target)
 
     elif isinstance(target, PrefetchDelayAt):
-        delay = _get_prefetch_item_delay(target)
-        delay.register_delay(target.time)
+        pass
+        # delay = _get_prefetch_item_delay(target)
+        # delay.register_delay(*target.delay_time)
 
 
 class DynamicsGroup(Module):

brainstate/nn/_elementwise.py

@@ -19,9 +19,8 @@ from typing import Optional
 
 import brainunit as u
 import jax.numpy as jnp
-import jax.typing
 
-from brainstate import random, functional as F
+from brainstate import functional as F
 from brainstate._state import ParamState
 from brainstate.typing import ArrayLike
 from ._module import ElementWiseBlock

brainstate/nn/_linear_mv.py

@@ -15,11 +15,11 @@
 
 from typing import Union, Callable, Optional
 
+import brainevent
 import brainunit as u
 import jax
 
 from brainstate import init
-import brainevent
 from brainstate._state import ParamState
 from brainstate.typing import Size, ArrayLike
 from ._module import Module

brainstate/nn/_module_test.py

@@ -15,172 +15,9 @@
 
 import unittest
 
-import jax.numpy as jnp
-
 import brainstate
 
 
-class TestDelay(unittest.TestCase):
-    def test_delay1(self):
-        a = brainstate.State(brainstate.random.random(10, 20))
-        delay = brainstate.nn.Delay(a.value)
-        delay.register_entry('a', 1.)
-        delay.register_entry('b', 2.)
-        delay.register_entry('c', None)
-
-        delay.init_state()
-        with self.assertRaises(KeyError):
-            delay.register_entry('c', 10.)
-
-    def test_rotation_delay(self):
-        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
-        t0 = 0.
-        t1, n1 = 1., 10
-        t2, n2 = 2., 20
-
-        rotation_delay.register_entry('a', t0)
-        rotation_delay.register_entry('b', t1)
-        rotation_delay.register_entry('c2', 1.9)
-        rotation_delay.register_entry('c', t2)
-
-        rotation_delay.init_state()
-
-        print()
-        # print(rotation_delay)
-        # print(rotation_delay.max_length)
-
-        for i in range(100):
-            brainstate.environ.set(i=i)
-            rotation_delay.update(jnp.ones((1,)) * i)
-            # print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c2'), rotation_delay.at('c'))
-            self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
-            self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
-            self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-
-    def test_concat_delay(self):
-        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
-        t0 = 0.
-        t1, n1 = 1., 10
-        t2, n2 = 2., 20
-
-        rotation_delay.register_entry('a', t0)
-        rotation_delay.register_entry('b', t1)
-        rotation_delay.register_entry('c', t2)
-
-        rotation_delay.init_state()
-
-        print()
-        for i in range(100):
-            brainstate.environ.set(i=i)
-            rotation_delay.update(jnp.ones((1,)) * i)
-            print(i, rotation_delay.at('a'), rotation_delay.at('b'), rotation_delay.at('c'))
-            self.assertTrue(jnp.allclose(rotation_delay.at('a'), jnp.ones((1,)) * i))
-            self.assertTrue(jnp.allclose(rotation_delay.at('b'), jnp.maximum(jnp.ones((1,)) * i - n1, 0.)))
-            self.assertTrue(jnp.allclose(rotation_delay.at('c'), jnp.maximum(jnp.ones((1,)) * i - n2, 0.)))
-        # brainstate.util.clear_buffer_memory()
-
-    def test_jit_erro(self):
-        rotation_delay = brainstate.nn.Delay(jnp.ones([1]), time=2., delay_method='concat', interp_method='round')
-        rotation_delay.init_state()
-
-        with brainstate.environ.context(i=0, t=0, jit_error_check=True):
-            rotation_delay.retrieve_at_time(-2.0)
-            with self.assertRaises(Exception):
-                rotation_delay.retrieve_at_time(-2.1)
-            rotation_delay.retrieve_at_time(-2.01)
-            with self.assertRaises(Exception):
-                rotation_delay.retrieve_at_time(-2.09)
-            with self.assertRaises(Exception):
-                rotation_delay.retrieve_at_time(0.1)
-            with self.assertRaises(Exception):
-                rotation_delay.retrieve_at_time(0.01)
-
-    def test_round_interp(self):
-        for shape in [(1,), (1, 1), (1, 1, 1)]:
-            for delay_method in ['rotation', 'concat']:
-                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                                     interp_method='round')
-                t0, n1 = 0.01, 0
-                t1, n1 = 1.04, 10
-                t2, n2 = 1.06, 11
-                rotation_delay.init_state()
-
-                @brainstate.compile.jit
-                def retrieve(td, i):
-                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
-                        return rotation_delay.retrieve_at_time(td)
-
-                print()
-                for i in range(100):
-                    t = i * brainstate.environ.get_dt()
-                    with brainstate.environ.context(i=i, t=t):
-                        rotation_delay.update(jnp.ones(shape) * i)
-                        print(i,
-                              retrieve(t - t0, i),
-                              retrieve(t - t1, i),
-                              retrieve(t - t2, i))
-                        self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.ones(shape) * i))
-                        self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
-                        self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
-
-    def test_linear_interp(self):
-        for shape in [(1,), (1, 1), (1, 1, 1)]:
-            for delay_method in ['rotation', 'concat']:
-                print(shape, delay_method)
-
-                rotation_delay = brainstate.nn.Delay(jnp.ones(shape), time=2., delay_method=delay_method,
-                                                     interp_method='linear_interp')
-                t0, n0 = 0.01, 0.1
-                t1, n1 = 1.04, 10.4
-                t2, n2 = 1.06, 10.6
-                rotation_delay.init_state()
-
-                @brainstate.compile.jit
-                def retrieve(td, i):
-                    with brainstate.environ.context(i=i, t=i * brainstate.environ.get_dt()):
-                        return rotation_delay.retrieve_at_time(td)
-
-                print()
-                for i in range(100):
-                    t = i * brainstate.environ.get_dt()
-                    with brainstate.environ.context(i=i, t=t):
-                        rotation_delay.update(jnp.ones(shape) * i)
-                        print(i,
-                              retrieve(t - t0, i),
-                              retrieve(t - t1, i),
-                              retrieve(t - t2, i))
-                        self.assertTrue(jnp.allclose(retrieve(t - t0, i), jnp.maximum(jnp.ones(shape) * i - n0, 0.)))
-                        self.assertTrue(jnp.allclose(retrieve(t - t1, i), jnp.maximum(jnp.ones(shape) * i - n1, 0.)))
-                        self.assertTrue(jnp.allclose(retrieve(t - t2, i), jnp.maximum(jnp.ones(shape) * i - n2, 0.)))
-
-    def test_rotation_and_concat_delay(self):
-        rotation_delay = brainstate.nn.Delay(jnp.ones((1,)))
-        concat_delay = brainstate.nn.Delay(jnp.ones([1]), delay_method='concat')
-        t0 = 0.
-        t1, n1 = 1., 10
-        t2, n2 = 2., 20
-
-        rotation_delay.register_entry('a', t0)
-        rotation_delay.register_entry('b', t1)
-        rotation_delay.register_entry('c', t2)
-        concat_delay.register_entry('a', t0)
-        concat_delay.register_entry('b', t1)
-        concat_delay.register_entry('c', t2)
-
-        rotation_delay.init_state()
-        concat_delay.init_state()
-
-        print()
-        for i in range(100):
-            brainstate.environ.set(i=i)
-            new = jnp.ones((1,)) * i
-            rotation_delay.update(new)
-            concat_delay.update(new)
-            self.assertTrue(jnp.allclose(rotation_delay.at('a'), concat_delay.at('a'), ))
-            self.assertTrue(jnp.allclose(rotation_delay.at('b'), concat_delay.at('b'), ))
-            self.assertTrue(jnp.allclose(rotation_delay.at('c'), concat_delay.at('c'), ))
-
-
 class TestModule(unittest.TestCase):
     def test_states(self):
         class A(brainstate.nn.Module):
@@ -201,8 +38,3 @@ class TestModule(unittest.TestCase):
         print(b.states())
         print(b.states(level=0))
         print(b.states(level=0))
-
-
-if __name__ == '__main__':
-    with brainstate.environ.context(dt=0.1):
-        unittest.main()

brainstate/nn/_synapse.py

@@ -23,7 +23,7 @@ import brainunit as u
 from brainstate import init, environ
 from brainstate._state import ShortTermState, HiddenState
 from brainstate.mixin import AlignPost
-from brainstate.typing import ArrayLike, Size, PyTree
+from brainstate.typing import ArrayLike, Size
 from ._dynamics import Dynamics
 from ._exp_euler import exp_euler_step
 

{brainstate-0.1.7.dist-info → brainstate-0.1.8.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.7
+Version: 0.1.8
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers

{brainstate-0.1.7.dist-info → brainstate-0.1.8.dist-info}/RECORD

@@ -1,4 +1,4 @@
-brainstate/__init__.py,sha256=l8qmh_ah9kUi2khglrj4O89ERt0tHdVT8WnVQu1lTcY,1496
+brainstate/__init__.py,sha256=bM6vJtv8DWVcxOUQ0GPpkCzmRkdTREdTutkSVv0OziQ,1496
 brainstate/_compatible_import.py,sha256=LUSZlA0APWozxM8Kf9pZrM2YbwY7X3jzVVHZInaBL7Y,4630
 brainstate/_state.py,sha256=PLzoYx13jIgnzyBxnktLTVPCFl6seG__aNIHS9A5Nms,60770
 brainstate/_state_test.py,sha256=b6uvZdVRyC4n6-fYzmHNry1b-gJ6zE_kRSxGinqiHaw,1638
@@ -31,8 +31,8 @@ brainstate/compile/_loop_collect_return.py,sha256=qOYGoD2TMZd2Px6y241GWNwSjOLyaK
 brainstate/compile/_loop_collect_return_test.py,sha256=pEJdcOthEM17q5kXYhgR6JfzzqibijD_O1VPXVe_Ml4,1808
 brainstate/compile/_loop_no_collection.py,sha256=tPkSxee41VexWEALpN1uuT78BDrX3uT1FHv8WZtov4c,7549
 brainstate/compile/_loop_no_collection_test.py,sha256=ivavF59xep_g9bV1SSdXd5E1f6nhc7EUBfcbgHpxbfg,1419
-brainstate/compile/_make_jaxpr.py,sha256=adQwsEOxFZQk8yr0SbB3Tjh5j3bBvZCZSh7JVf6ymCA,37387
-brainstate/compile/_make_jaxpr_test.py,sha256=xPusEJikMQRfoOmUFrugGcE5UpRm0giHoL_NPomN5rI,4791
+brainstate/compile/_make_jaxpr.py,sha256=BOJa3H1Fa0Esh3Al3CDaS351w45vk8xPDc9R_nk8jh8,37367
+brainstate/compile/_make_jaxpr_test.py,sha256=_aRwYt932dXlcntFNYlm5asRRE97oGp4F8hrgNO63SY,5164
 brainstate/compile/_progress_bar.py,sha256=AhAyyI_ckzgaj0PSj1ep1hq8rGzQLSyC2aVCYaT-e-o,7502
 brainstate/compile/_unvmap.py,sha256=9S42MeTmFJa8nfBI_AjEfrAdUsDmM3KFat59O8zXIEw,4120
 brainstate/compile/_util.py,sha256=QD7lvS4Zb3P2HPNIm6mG8Rl_Lk2tQj0znOCPcH95XnI,6275
@@ -60,12 +60,13 @@ brainstate/nn/_collective_ops_test.py,sha256=bwq0DApcsk0_2xpxMl0_e2cGKT63g5rSngp
 brainstate/nn/_common.py,sha256=qHAOID_eeKiPUXk_ION65sbYXyF-ddH5w5BayvH8Thg,6431
 brainstate/nn/_conv.py,sha256=Zk-yj34n6CkjntcM9xpMGLTxKNfWdIWsTsoGbtdL0yU,18448
 brainstate/nn/_conv_test.py,sha256=2lcUTG7twkyhuyKwuBux-NgU8NU_W4Cp1-G8EyDJ_uk,8862
-brainstate/nn/_delay.py,sha256=l36FBgNhfL64tM3VGOsJNTtKr44HjxxtBWMFFCm3Pks,17361
+brainstate/nn/_delay.py,sha256=_FFKQ09-NHLyYPcfuvoHzpL2zc9Twu_p3k33VDbgxKo,19523
+brainstate/nn/_delay_test.py,sha256=O5eY-LTuzxLi6hjsC-lC12mhYe8xL36cWzvEwnG7h4g,8187
 brainstate/nn/_dropout.py,sha256=Dq3hQrOBT6gODlDbcoag6zLGXTU_p_2MhnmVsV57Hds,17783
 brainstate/nn/_dropout_test.py,sha256=L46PvC2OA7EnS4MsRhh_YnvKheHYNafOsKM8uzux_zo,4446
-brainstate/nn/_dynamics.py,sha256=PTTGDqe-uRvrTZbhNGy1O4hT_wG4vxbOXyfbOx2itaw,48341
+brainstate/nn/_dynamics.py,sha256=QVCQqdglR1p-vMFzXvvEWfK1ltg4PJhUxSEEbbFP95s,48236
 brainstate/nn/_dynamics_test.py,sha256=w7AV57LdhbBNYprdFpKq8MFSCbXKVkGgp_NbL3ANX3I,2769
-brainstate/nn/_elementwise.py,sha256=l3Wk65UDLFc3xFH1jwTXRwKAVmt3dUeG0iF_Al0w1_0,33487
+brainstate/nn/_elementwise.py,sha256=_9O4S-gBZiTmsx5lJ3WdQVLcACBmXTwPTY8pnUoK7KY,33460
 brainstate/nn/_elementwise_test.py,sha256=_dd9eX2ZJ7p24ahuoapCaRTZ0g1boufXMyqHFx1d4WY,5688
 brainstate/nn/_embedding.py,sha256=SaAJbgXmuJ8XlCOX9ob4yvmgh9Fk627wMguRzJMJ1H8,2138
 brainstate/nn/_exp_euler.py,sha256=WTpZm-XQmsdMLNazY7wIu8eeO6pK0kRzt2lJnhEgMIk,3293
@@ -74,12 +75,12 @@ brainstate/nn/_fixedprob.py,sha256=KGXohiU0wZnFIQDuwiRUTFsbsr8R0p8zgi5UZDuv1Bk,1
 brainstate/nn/_fixedprob_test.py,sha256=qbRBh-MpMtEOsg492gFu2w9-FOP9z_bXapm-Q0gLLYM,3929
 brainstate/nn/_inputs.py,sha256=hMPkx9qDBpJWPshZXLF4H1QiYK1-46wntHUIlG7cT7c,20603
 brainstate/nn/_linear.py,sha256=FnPxATdT66DecjTW0tUTyL6clQmN_cG8kPC3qDWUE6A,14500
-brainstate/nn/_linear_mv.py,sha256=6hDXx4yPqRSa7uIsW9f9eJuy23dcXN9Mp2_lSvw8BDA,2635
+brainstate/nn/_linear_mv.py,sha256=DyJ0OOHcvuJrBvBhV0CbmVgw-6BCg_ms-Rso9eB5e1E,2635
 brainstate/nn/_linear_mv_test.py,sha256=ZCM1Zy6mImQfCfdZOGnTwkiLLPXK5yalv1Ts9sWZuPA,3864
 brainstate/nn/_linear_test.py,sha256=eIS-VCR3QmXB_byO1Uexg65Pv48CBRUA_Je-UGrFVTY,2925
 brainstate/nn/_ltp.py,sha256=_najNUyfaFYcOlUTm7ThJopInbos3kwJyrm-QUfI-hc,861
 brainstate/nn/_module.py,sha256=jlFaoltT2F8a7cxsyu6fXp7mfIkEsk-JD3G1Xh3Ay8I,12783
-brainstate/nn/_module_test.py,sha256=sjT7t-N4ZZKPN_MujNv8bcT5uZQD-CmU5nKz9KwlBUc,8963
+brainstate/nn/_module_test.py,sha256=v4ZHwjvrEaqHzCgyR8MKIB-iNhdtaCLOc55EORXTjPg,1451
 brainstate/nn/_neuron.py,sha256=2walTScvL034LS53pArDASXz6z26SSPbmCvchWWjkUU,27441
 brainstate/nn/_neuron_test.py,sha256=QF8pixUqA5Oj7MrNi2NR8VAnfGpAvNpwV2mBc3e_pTY,6393
 brainstate/nn/_normalizations.py,sha256=YSC1W7JaexoZ8tKcy5B-dm-_x8GuPzS_6XBfkaKpdXM,37464
@@ -92,7 +93,7 @@ brainstate/nn/_rate_rnns_test.py,sha256=__hhx7e6LX_1mDLLQyIi4TNCaFAWnOVSTIgwHNjz
 brainstate/nn/_readout.py,sha256=OJjSba5Wr7dtUXqYhAv1D7BUGOI-lAmg6urxPBrZe3c,7116
 brainstate/nn/_readout_test.py,sha256=L2T0-SkiACxkY_I5Pbnbmy0Zw3tbpV3l5xVzAw42f2g,2136
 brainstate/nn/_stp.py,sha256=-ahDEqSp8bQsU_nUK4jks8fjMYKgIbO0v7zpyGVuXtA,8645
-brainstate/nn/_synapse.py,sha256=E43DxH0PW3pa6rKn9UWkD_PxIdpMAo04Yy41gaR3KPc,19868
+brainstate/nn/_synapse.py,sha256=I7npO9rg6bvzCiCNLiTqEOxUjGaDVlkOgmkT9M1V4xY,19860
 brainstate/nn/_synapse_test.py,sha256=xmCWFxZUIM2YtmW5otKnADGCCK__4JpXmSYcZ3wzlQM,4994
 brainstate/nn/_synaptic_projection.py,sha256=UFgzsMB1VZ9ieumwmaTIC7irLrZ4pmwfiuOYomkwXG4,17917
 brainstate/nn/_synouts.py,sha256=jWQP1-qXFpdYgyUSJNFD7_bk4_-67ok36br-OzbcSXY,4524
@@ -124,8 +125,8 @@ brainstate/util/pretty_repr.py,sha256=7Xp7IFNUeP7cGlpvwwJyBslbQVnXEqC1I6neV1Jx1S
 brainstate/util/pretty_table.py,sha256=uJVaamFGQ4nKP8TkEGPWXHpzjMecDo2q1Ah6XtRjdPY,108117
 brainstate/util/scaling.py,sha256=U6DM-afPrLejiGqo1Nla7z4YbTBVicctsBEweurr_mk,7524
 brainstate/util/struct.py,sha256=2Y_wuDFQ6ldl_H4_w0IjzAtkbHooVgdsVbnT7Z6_Efc,17528
-brainstate-0.1.7.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.7.dist-info/METADATA,sha256=PrV7tP9tHHf5d8ZOUoXPa9t5Sm0WBo41thLQVjfgHuA,4122
-brainstate-0.1.7.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
-brainstate-0.1.7.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.7.dist-info/RECORD,,
+brainstate-0.1.8.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.8.dist-info/METADATA,sha256=85Nudry-TozOxctIhW8zx7DUgOxPVmML3XCq7YqYDFs,4122
+brainstate-0.1.8.dist-info/WHEEL,sha256=AHX6tWk3qWuce7vKLrj7lnulVHEdWoltgauo8bgCXgU,109
+brainstate-0.1.8.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.8.dist-info/RECORD,,