brainstate 0.1.0.post20250218__py2.py3-none-any.whl → 0.1.0.post20250315__py2.py3-none-any.whl

brainstate/augment/_autograd.py

@@ -29,11 +29,12 @@ The wrapped gradient transformations here are made possible by using the followi
 
 from __future__ import annotations
 
-import brainunit as u
-import jax
 from functools import wraps, partial
 from typing import Union, Callable, Dict, Sequence, Optional, Any, Tuple, TypeVar, Iterator
 
+import brainunit as u
+import jax
+
 from brainstate._state import State
 from brainstate._utils import set_module_as
 from brainstate.compile._make_jaxpr import StatefulFunction
@@ -195,6 +196,7 @@ class GradientTransform(PrettyRepr):
             grad_states = {k: v for k, v in grad_states.items()}
         self._grad_states, self._grad_tree = jax.tree.flatten(grad_states)
         self._grad_state_ids = [id(v) for v in self._grad_states]
+        self._grad_id_to_state = {id(v): v for v in self._grad_states}
         if any(not isinstance(v, State) for v in self._grad_states):
             raise TypeError("All grad_states must be State instances.")
 
@@ -250,12 +252,20 @@ class GradientTransform(PrettyRepr):
         """
         grad_vals = dict()
         other_vals = dict()
+        all_ids = set(self._grad_state_ids)
         for st in state_trace.states:
             id_ = id(st)
-            if id_ in self._grad_state_ids:
+            if id_ in all_ids:
                 grad_vals[id_] = st.value
+                all_ids.remove(id_)
             else:
                 other_vals[id_] = st.value
+        if len(all_ids):
+            err = f"Some states are not found in the state trace when performing gradient transformations.\n "
+            for i, id_ in enumerate(all_ids):
+                st = self._grad_id_to_state[id_]
+                st.raise_error_with_source_info(ValueError(err + str(st)))
+
         return grad_vals, other_vals
 
     def _merge_state_vals(self, grad_vals: Dict, other_vals: Dict, state_trace):

brainstate/augment/_mapping.py

@@ -16,8 +16,6 @@
 from __future__ import annotations
 
 import functools
-import jax
-from jax.interpreters.batching import BatchTracer
 from typing import (
     Any,
     TypeVar,
@@ -32,6 +30,9 @@ from typing import (
     List
 )
 
+import jax
+from jax.interpreters.batching import BatchTracer
+
 from brainstate._state import State, catch_new_states
 from brainstate.compile import scan, StatefulFunction
 from brainstate.random import RandomState, DEFAULT
@@ -378,8 +379,15 @@ def _vmap_transform(
         # call the function
         return f(*args)
 
+    def _set_axis_env(batch_size):
+        axis_env = None if axis_name is None else [(axis_name, batch_size)]
+        stateful_fn.axis_env = axis_env
+
     # stateful function
-    stateful_fn = StatefulFunction(_vmap_fn_for_compilation, name='vmap')
+    stateful_fn = StatefulFunction(
+        _vmap_fn_for_compilation,
+        name='vmap',
+    )
 
     @functools.wraps(f)
     def new_fn_for_vmap(
@@ -506,6 +514,10 @@ def _vmap_transform(
             st_in_axes = 0
 
         # compile stateful function
+        batch_size = None
+        if axis_name is not None:
+            batch_size = _get_batch_size(args, in_axes, axis_to_in_states, axis_size)
+            _set_axis_env(batch_size)
         cache_key = _compile_stateful_function(
             stateful_fn,
             (st_in_axes, in_axes),
@@ -518,7 +530,8 @@ def _vmap_transform(
         rng_sets = set(rngs)
         if len(rngs):
             # batch size
-            batch_size = _get_batch_size(args, in_axes, axis_to_in_states, axis_size)
+            if batch_size is None:
+                batch_size = _get_batch_size(args, in_axes, axis_to_in_states, axis_size)
             rng_keys = tuple(rng.split_key(batch_size) for rng in rngs)
             rng_backup = tuple(rng.split_key() for rng in rngs)
         else:
@@ -905,11 +918,11 @@ def map(
         g = lambda _, x: ((), vmap(f)(*x))
         _, scan_ys = scan(g, (), scan_xs)
         if remainder_xs is None:
-            ys = jax.tree.map(lambda x: flatten_(x), scan_ys)
+            ys = jax.tree.map(lambda x: _flatten(x), scan_ys)
         else:
             remainder_ys = vmap(f)(*remainder_xs)
             ys = jax.tree.map(
-                lambda x, y: jax.lax.concatenate([flatten_(x), y], dimension=0),
+                lambda x, y: jax.lax.concatenate([_flatten(x), y], dimension=0),
                 scan_ys,
                 remainder_ys,
             )
@@ -919,7 +932,7 @@ def map(
     return ys
 
 
-def flatten_(x):
+def _flatten(x):
     return x.reshape(-1, *x.shape[2:])
 
 
@@ -935,17 +948,19 @@ def _vmap_new_states_transform(
     # -- brainstate specific arguments -- #
     state_tag: str | None = None,
     state_to_exclude: Filter | None = None,
+    in_states: Dict[int, Dict] | Any | None = None,
+    out_states: Dict[int, Dict] | Any | None = None,
 ):
-
     # TODO: How about nested call ``vmap_new_states``?
 
-
     @vmap(
         in_axes=in_axes,
         out_axes=out_axes,
         axis_name=axis_name,
         axis_size=axis_size,
         spmd_axis_name=spmd_axis_name,
+        in_states=in_states,
+        out_states=out_states,
     )
     def new_fun(args):
         # call the function
@@ -988,6 +1003,8 @@ def vmap_new_states(
     # -- brainstate specific arguments -- #
     state_tag: str | None = None,
     state_to_exclude: Filter = None,
+    in_states: Dict[int, Dict] | Any | None = None,
+    out_states: Dict[int, Dict] | Any | None = None,
 ):
     """
     Vectorize a function over new states created within it.
@@ -1019,6 +1036,8 @@ def vmap_new_states(
             spmd_axis_name=spmd_axis_name,
             state_tag=state_tag,
             state_to_exclude=state_to_exclude,
+            in_states=in_states,
+            out_states=out_states,
         )
     else:
         return _vmap_new_states_transform(
@@ -1030,4 +1049,6 @@ def vmap_new_states(
             spmd_axis_name=spmd_axis_name,
             state_tag=state_tag,
             state_to_exclude=state_to_exclude,
+            in_states=in_states,
+            out_states=out_states,
         )

brainstate/augment/_mapping_test.py

@@ -15,12 +15,14 @@
 
 from __future__ import annotations
 
+import unittest
+
 import jax
 import jax.numpy as jnp
 import numpy as np
-import unittest
 
 import brainstate as bst
+import brainstate.augment
 from brainstate.augment._mapping import BatchAxisError
 from brainstate.augment._mapping import _remove_axis
 
@@ -264,6 +266,55 @@ class TestVmap(unittest.TestCase):
         res2 = jax.vmap(f, axis_size=10)()
         self.assertTrue(jnp.all((res2[0] == res2[1:])))
 
+    def test_axis(self):
+        def f(x):
+            return x - jax.lax.pmean(x, 'i')
+
+        r = jax.vmap(f, axis_name='i')(jnp.arange(10))
+        print(r)
+
+        r2 = brainstate.augment.vmap(f, axis_name='i')(jnp.arange(10))
+        print(r2)
+        self.assertTrue(jnp.allclose(r, r2))
+
+    def test_vmap_init(self):
+        class Foo(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.a = bst.ParamState(jnp.arange(4))
+                self.b = bst.ShortTermState(jnp.arange(4))
+
+            def init_state_v1(self, *args, **kwargs):
+                self.c = bst.State(jnp.arange(4))
+
+            def init_state_v2(self):
+                self.d = bst.State(self.c.value * 2.)
+
+        foo = Foo()
+
+        @brainstate.augment.vmap_new_states(state_tag='new1', axis_size=5)
+        def init1():
+            foo.init_state_v1()
+
+        init1()
+        print(foo.c.value)
+
+        @brainstate.augment.vmap_new_states(state_tag='new2', axis_size=5, in_states=foo.states('new1'))
+        def init2():
+            foo.init_state_v2()
+
+        init2()
+        print(foo.c.value)
+        print(foo.d.value)
+
+        self.assertTrue(
+            jnp.allclose(
+                foo.d.value,
+                foo.c.value * 2.
+            )
+        )
+
+
 
 class TestMap(unittest.TestCase):
     def test_map(self):

brainstate/compile/_make_jaxpr.py

@@ -753,8 +753,8 @@ def _make_jaxpr(
         if jax.__version_info__ < (0, 5, 0):
             debug_info = pe.debug_info(fun, in_tree, out_tree, True, 'make_jaxpr')
         with ExitStack() as stack:
-            for axis_name, size in axis_env or []:
-                stack.enter_context(jax.core.extend_axis_env(axis_name, size, None))
+            if axis_env is not None:
+                stack.enter_context(jax.core.extend_axis_env_nd(axis_env))
             if jax.__version_info__ < (0, 5, 0):
                 jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
             else:

brainstate/functional/_spikes.py

@@ -159,7 +159,7 @@ def spike_bitwise_ixor(x, y):
 
 
 def spike_bitwise(x, y, op: str):
-    """
+    r"""
     Perform bitwise operations on spike tensors.
 
     This function applies various bitwise operations on spike tensors based on the specified operation.

brainstate/nn/_collective_ops.py

@@ -18,9 +18,7 @@ from __future__ import annotations
 from collections import namedtuple
 
 import jax
-from typing import (
-    Callable, TypeVar, Tuple, Any, Dict
-)
+from typing import Callable, TypeVar, Tuple, Any, Dict
 
 from brainstate._state import catch_new_states
 from brainstate._utils import set_module_as
@@ -103,7 +101,7 @@ def call_all_functions(
     on each node. It respects the call order of functions if defined, and provides options for
     handling cases where the specified function does not exist on a node.
 
-    Parameters:
+    Parameters
     -----------
     target : T
         The target module on which to call functions.
@@ -121,12 +119,12 @@ def call_all_functions(
         - 'raise': Raise an exception (default)
         - 'pass' or 'none': Skip the node and continue
         
-    Returns:
+    Returns
     --------
     T
         The target module after calling the specified function on all applicable nodes.
 
-    Raises:
+    Raises
     -------
     AssertionError
         If fun_name is not a string or kwargs is not a dictionary.
@@ -186,7 +184,7 @@ def vmap_call_all_functions(
     This function vectorizes the process of calling a specified function across multiple instances
     of the target module, effectively batching the operation.
 
-    Parameters:
+    Parameters
     -----------
     target : T
         The target module on which to call functions.
@@ -208,12 +206,12 @@ def vmap_call_all_functions(
         - 'raise': Raise an exception (default)
         - 'pass' or 'none': Skip the node and continue
 
-    Returns:
+    Returns
     --------
     T
         The target module after applying the vectorized function call on all applicable nodes.
 
-    Raises:
+    Raises
     -------
     AssertionError
         If axis_size is not specified or is not a positive integer.
@@ -304,7 +302,7 @@ def vmap_init_all_states(
     This function applies vectorized mapping (vmap) to initialize states across multiple
     instances of the target module, effectively batching the initialization process.
 
-    Parameters:
+    Parameters
     -----------
     target : T
         The target module whose states are to be initialized.
@@ -319,12 +317,12 @@ def vmap_init_all_states(
     state_tag : str | None, optional
         A tag to be used for catching new states.
 
-    Returns:
+    Returns
     --------
     T
         The target module with initialized states.
 
-    Raises:
+    Raises
     -------
     AssertionError
         If axis_size is not specified or is not greater than 0.
@@ -413,7 +411,7 @@ def vmap_reset_all_states(
     This function applies vectorized mapping (vmap) to reset states across multiple
     instances of the target module, effectively batching the reset process.
 
-    Parameters:
+    Parameters
     -----------
     target : T
         The target module whose states are to be reset.
@@ -428,12 +426,12 @@ def vmap_reset_all_states(
     tag : str | None, optional
         A tag to be used for catching new states.
 
-    Returns:
+    Returns
     --------
     T
         The target module with reset states.
 
-    Raises:
+    Raises
     -------
     AssertionError
         If axis_size is not specified or is not greater than 0.
@@ -486,7 +484,7 @@ def save_all_states(target: Module, **kwargs) -> Dict:
     Args:
       target: Module. The node to save its states.
 
-    Returns:
+    Returns
       Dict. The state dict for serialization.
     """
     return {key: node.save_state(**kwargs) for key, node in target.nodes().items()}

brainstate/nn/_dyn_impl/_dynamics_synapse.py

@@ -17,8 +17,9 @@
 
 from __future__ import annotations
 
+from typing import Optional, Callable
+
 import brainunit as u
-from typing import Optional
 
 from brainstate import init, environ
 from brainstate._state import ShortTermState, HiddenState
@@ -54,7 +55,7 @@ class Expon(Synapse, AlignPost):
         in_size: Size,
         name: Optional[str] = None,
         tau: ArrayLike = 8.0 * u.ms,
-        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(unit=u.mS),
     ):
         super().__init__(name=name, in_size=in_size)
 
@@ -85,7 +86,7 @@ class DualExpon(Synapse, AlignPost):
         tau_decay: ArrayLike = 10.0 * u.ms,
         tau_rise: ArrayLike = 1.0 * u.ms,
         A: Optional[ArrayLike] = None,
-        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(unit=u.mS),
     ):
         super().__init__(name=name, in_size=in_size)
 
@@ -133,7 +134,7 @@ class Alpha(Synapse):
         in_size: Size,
         name: Optional[str] = None,
         tau: ArrayLike = 8.0 * u.ms,
-        g_initializer: ArrayLike = init.ZeroInit(unit=u.mS),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(unit=u.mS),
     ):
         super().__init__(name=name, in_size=in_size)
 
@@ -321,7 +322,7 @@ class AMPA(Synapse):
         beta: ArrayLike = 0.18 / u.ms,
         T: ArrayLike = 0.5 * u.mM,
         T_dur: ArrayLike = 0.5 * u.ms,
-        g_initializer: ArrayLike = init.ZeroInit(),
+        g_initializer: ArrayLike | Callable = init.ZeroInit(),
     ):
         super().__init__(name=name, in_size=in_size)
 
@@ -394,5 +395,14 @@ class GABAa(AMPA):
         beta: ArrayLike = 0.18 / u.ms,
         T: ArrayLike = 1.0 * u.mM,
         T_dur: ArrayLike = 1.0 * u.ms,
+        g_initializer: ArrayLike | Callable = init.ZeroInit(),
     ):
-        super().__init__(alpha=alpha, beta=beta, T=T, T_dur=T_dur, name=name, in_size=in_size)
+        super().__init__(
+            alpha=alpha,
+            beta=beta,
+            T=T,
+            T_dur=T_dur,
+            name=name,
+            in_size=in_size,
+            g_initializer=g_initializer
+        )

brainstate/nn/_exp_euler.py

@@ -49,13 +49,13 @@ def exp_euler_step(
     should have units of ( [X]/\sqrt{[T]} ).
 
     Args:
-      fun: Callable. The function to be solved.
-      diffusion: Callable. The diffusion function.
-      *args: The input arguments.
-      drift: Callable. The drift function.
+        fun: Callable. The function to be solved.
+        diffusion: Callable. The diffusion function.
+        *args: The input arguments.
+        drift: Callable. The drift function.
 
     Returns:
-      The one-step solution of the ODE.
+        The one-step solution of the ODE.
     """
     assert callable(fn), 'The input function should be callable.'
     assert len(args) > 0, 'The input arguments should not be empty.'

brainstate/optim/_sgd_optimizer.py

@@ -201,9 +201,11 @@ class SGD(_WeightDecayOptimizer):
     def update(self, grads: dict):
         lr = self.lr()
         weight_values, grad_values = to_same_dict_tree(self.param_states, grads)
-        updates = jax.tree.map(functools.partial(_sgd, lr=lr, weight_decay=self.weight_decay),
-                               weight_values,
-                               grad_values)
+        updates = jax.tree.map(
+            functools.partial(_sgd, lr=lr, weight_decay=self.weight_decay),
+            weight_values,
+            grad_values
+        )
         self.param_states.assign_values(updates)
         self.lr.step_call()
 
@@ -324,12 +326,16 @@ class MomentumNesterov(_WeightDecayOptimizer):
     def update(self, grads: dict):
         lr = self.lr()
         states_values, grad_values, momentum_values = to_same_dict_tree(self.param_states, grads, self.momentum_states)
-        momentum_values = jax.tree.map(lambda mv, gv: self.momentum * mv - lr * gv,
-                                       momentum_values,
-                                       grad_values)
-        weight_values = jax.tree.map(functools.partial(_sgd, lr=lr, weight_decay=self.weight_decay),
-                                     states_values,
-                                     momentum_values)
+        momentum_values = jax.tree.map(
+            lambda mv, gv: self.momentum * mv - lr * gv,
+            momentum_values,
+            grad_values
+        )
+        weight_values = jax.tree.map(
+            functools.partial(_sgd, lr=lr, weight_decay=self.weight_decay),
+            states_values,
+            momentum_values
+        )
         self.param_states.assign_values(weight_values)
         self.momentum_states.assign_values(momentum_values)
         self.lr.step_call()
@@ -388,11 +394,21 @@ class Adagrad(_WeightDecayOptimizer):
     def update(self, grads: dict):
         lr = self.lr()
         cache_values, grad_values, weight_values = to_same_dict_tree(self.cache_states, grads, self.param_states)
-        cache_values = jax.tree.map(lambda cv, gv: cv + gv ** 2, cache_values, grad_values)
-        updates = jax.tree.map(lambda cv, gv: lr * gv / jnp.sqrt(cv + self.epsilon), cache_values, grad_values)
-        weight_values = jax.tree.map(functools.partial(_sgd, weight_decay=self.weight_decay),
-                                     weight_values,
-                                     updates)
+        cache_values = jax.tree.map(
+            lambda cv, gv: cv + gv ** 2,
+            cache_values,
+            grad_values
+        )
+        updates = jax.tree.map(
+            lambda cv, gv: lr * gv / jnp.sqrt(cv + self.epsilon),
+            cache_values,
+            grad_values
+        )
+        weight_values = jax.tree.map(
+            functools.partial(_sgd, weight_decay=self.weight_decay),
+            weight_values,
+            updates
+        )
         self.cache_states.assign_values(cache_values)
         self.param_states.assign_values(weight_values)
         self.lr.step_call()
@@ -605,13 +621,28 @@ class Adam(_WeightDecayOptimizer):
         lr = lr / (1 - self.beta1 ** (self.lr.last_epoch.value + 2))
         lr = lr * jnp.sqrt(1 - self.beta2 ** (self.lr.last_epoch.value + 2))
         weight_values, grad_values, m1_values, m2_values = to_same_dict_tree(
-            self.param_states, grads, self.m1_states, self.m2_states)
-        m1_values = jax.tree.map(lambda m1, gv: self.beta1 * m1 + (1 - self.beta1) * gv, m1_values, grad_values)
-        m2_values = jax.tree.map(lambda m2, gv: self.beta2 * m2 + (1 - self.beta2) * gv ** 2, m2_values, grad_values)
-        update = jax.tree.map(lambda m1, m2: lr * m1 / (jnp.sqrt(m2) + self.eps), m1_values, m2_values)
-        weight_values = jax.tree.map(functools.partial(_sgd, weight_decay=self.weight_decay),
-                                     weight_values,
-                                     update)
+            self.param_states, grads, self.m1_states, self.m2_states
+        )
+        m1_values = jax.tree.map(
+            lambda m1, gv: self.beta1 * m1 + (1 - self.beta1) * gv,
+            m1_values,
+            grad_values
+        )
+        m2_values = jax.tree.map(
+            lambda m2, gv: self.beta2 * m2 + (1 - self.beta2) * gv ** 2,
+            m2_values,
+            grad_values
+        )
+        update = jax.tree.map(
+            lambda m1, m2: lr * m1 / (jnp.sqrt(m2) + self.eps),
+            m1_values,
+            m2_values
+        )
+        weight_values = jax.tree.map(
+            functools.partial(_sgd, weight_decay=self.weight_decay),
+            weight_values,
+            update
+        )
         self.param_states.assign_values(weight_values)
         self.m1_states.assign_values(m1_values)
         self.m2_states.assign_values(m2_values)

{brainstate-0.1.0.post20250218.dist-info → brainstate-0.1.0.post20250315.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250218
+Version: 0.1.0.post20250315
 Summary: A ``State``-based Transformation System for Program Compilation and Augmentation.
 Home-page: https://github.com/chaobrain/brainstate
 Author: BrainState Developers
@@ -62,6 +62,7 @@ Requires-Dist: jaxlib[tpu] ; extra == 'tpu'
     <a href="https://badge.fury.io/py/brainstate"><img alt="PyPI version" src="https://badge.fury.io/py/brainstate.svg"></a>
     <a href="https://github.com/chaobrain/brainstate/actions/workflows/CI.yml"><img alt="Continuous Integration" src="https://github.com/chaobrain/brainstate/actions/workflows/CI.yml/badge.svg"></a>
     <a href="https://pepy.tech/projects/brainstate"><img src="https://static.pepy.tech/badge/brainstate" alt="PyPI Downloads"></a>
+    <a href="https://doi.org/10.5281/zenodo.14970015"><img src="https://zenodo.org/badge/811300394.svg" alt="DOI"></a>
 </p>
 
 
@@ -81,8 +82,8 @@ The official documentation is hosted on Read the Docs: [https://brainstate.readt
 
 
 
-## See also the BDP ecosystem
+## See also the brain modeling ecosystem
 
-We are building the Brain Dynamics Programming ecosystem: https://ecosystem-for-brain-dynamics.readthedocs.io/
+We are building the brain modeling ecosystem: https://brainmodeling.readthedocs.io/
 
 

{brainstate-0.1.0.post20250218.dist-info → brainstate-0.1.0.post20250315.dist-info}/RECORD

@@ -10,12 +10,12 @@ brainstate/surrogate.py,sha256=wWYw-TxaFxHVneXuHjWD1UtTcOTk3XRSnhRtUkt_Hb8,53580
 brainstate/transform.py,sha256=vZWzO4F7qTsXL_SiVQPlTz0l9b_hRo9D-igETfgCTy0,758
 brainstate/typing.py,sha256=988gX1tvwtyYnYjmej90OaRxoMoBIPO0-DSrXXGxojM,10523
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
-brainstate/augment/_autograd.py,sha256=eAXdFPOEei2yTHtrhMXQBpC-36fmc6YZu3YWw4q_Wmk,29682
+brainstate/augment/_autograd.py,sha256=hfDoa2HbkRn-InOS0yOcb6gEZ2DLNqtWA133P8-hvIo,30138
 brainstate/augment/_autograd_test.py,sha256=2wCC8aUcDp2IHgF7wr1GK5HwWfELXni5PpA-082azuU,44058
 brainstate/augment/_eval_shape.py,sha256=jgsS197Nizehr9A2nGaQPE7NuNujhFhmR3J96hTicX8,3890
 brainstate/augment/_eval_shape_test.py,sha256=LFOJx7CWltmRLXdGY175UebLwtEMz2CzJ_gLqMZsJTw,1393
-brainstate/augment/_mapping.py,sha256=XqKRM88yAvvhWzSUqwq67sX49xV8oYtUdcmQ0KO9hVs,42532
-brainstate/augment/_mapping_test.py,sha256=V9NOpiZydDDUUYnLLCJVU9KJ-2IRNuVbL-ziuW-pU1I,12033
+brainstate/augment/_mapping.py,sha256=BPwpD7jX4xRNl4BdAsKGoF45MKbmEF9Lyyp11pJucIg,43356
+brainstate/augment/_mapping_test.py,sha256=-4HJXmJw_6SD9dQnHTBjgYVuq6VTVjz0xpc9v2CJVNw,13414
 brainstate/augment/_random.py,sha256=ikRzNoDDE2BkARajDsBhNlngCUrghzGSZUDmEGvVors,5386
 brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=3wv-f89oo94XeWwRV5LcRot0Nz7xTk5_PdjEDyUMsoo,9394
@@ -30,7 +30,7 @@ brainstate/compile/_loop_collect_return.py,sha256=-LsP7fkHmAyGnDOKa3BxxYOEWe8M2J
 brainstate/compile/_loop_collect_return_test.py,sha256=D9RQ5RyQHkqBr4nmSK-yM_uge3EC6uVm_Dzy42g3vtg,1802
 brainstate/compile/_loop_no_collection.py,sha256=2OEVtv5XztOx-e0focZ1UnWkXmFzmDskjHJXuVXmuhA,7587
 brainstate/compile/_loop_no_collection_test.py,sha256=oStB1CSG_iLp9sHdXd1hJNFvlxbzjck9Iy4sABoJDj4,1419
-brainstate/compile/_make_jaxpr.py,sha256=QGbVU_Y6pEfFgr61v_gmsRS3HXcHp7ILV_JZk_e3J4o,33213
+brainstate/compile/_make_jaxpr.py,sha256=8iV8XyvkMH3n3wbEWZAgZtbrUxryljwQJD6o5DMW9Lc,33189
 brainstate/compile/_make_jaxpr_test.py,sha256=fZe3K4RHFLmMAeXZoFZ5RyxgXvncTcuMQdjmOROJtKU,4365
 brainstate/compile/_progress_bar.py,sha256=3Z3OVcc5sl9FK9Fkt813l20MNzEfa6UZ9lJrvSgXTCU,7522
 brainstate/compile/_unvmap.py,sha256=uCvQjvb8J7kT0kalX576mrAPvQuCh_W76EPdgZ53kTM,4230
@@ -40,7 +40,7 @@ brainstate/functional/_activations.py,sha256=VmCU9HOKWbysxuJFBN-JsShS4loNMG_E6IX
 brainstate/functional/_activations_test.py,sha256=-bCijTvo4Wo_P283RYKYMPcTLsjhu5i2X9ySdf1ayEY,13034
 brainstate/functional/_normalization.py,sha256=L3S4DIF1EztrlE4_KHX7j_m6Mw0mpAwnx5UTAX6YYBU,2603
 brainstate/functional/_others.py,sha256=eBV43WqQsDvHkkwX0xbqCRoIoJlngMFLSUKgleH2dt0,1735
-brainstate/functional/_spikes.py,sha256=EZqwJctEElHhxaC9tsm8WKkXTQJJzo901Db0I3QS3iM,7086
+brainstate/functional/_spikes.py,sha256=7FTfCfEN1mjlY-EULzCisk7_NOmxZPj-mp-ODncW7R0,7087
 brainstate/graph/__init__.py,sha256=noo4TjBg6iEhjjwk0sAGUhR7Ge-z8Vnc2rLYUvnqttw,1295
 brainstate/graph/_graph_node.py,sha256=JE1Tc0mK3nJFWUFzXE53MWWiEYiXJO5VdqZEYKbXlw0,6872
 brainstate/graph/_graph_node_test.py,sha256=BFGfdzZFDHI0XK7hHotSVWKt3em1taGvn8FHF9NCXx8,2702
@@ -54,10 +54,10 @@ brainstate/init/_random_inits_test.py,sha256=lBL2RQdBSZ88Zqz4IMdbHJMvDi7ooZq6caC
 brainstate/init/_regular_inits.py,sha256=DmVMajugfyYFNUMzgFdDKMvbBu9hMWxkfDd-50uhoLg,3187
 brainstate/init/_regular_inits_test.py,sha256=tJl4aOkclllJIfKzJTbc0cfYCw2SoBsx8_G123RnqbU,1842
 brainstate/nn/__init__.py,sha256=ar1hDUYbSO6oadMpbuS9FWZvZB_iyFzM8CwMK-RNDzM,1823
-brainstate/nn/_collective_ops.py,sha256=yQNBnh-XVEFnTg-Ga14mHOCGtGxiTkL9MYKdNjJF1BI,17535
+brainstate/nn/_collective_ops.py,sha256=NI9BT-908TbIlXLMjbWsPyI5YLZD_cCkSKGeOY-qO60,17512
 brainstate/nn/_collective_ops_test.py,sha256=yW7NNYsGFglFRFkqVlpGSY6WLnU-h8GlK6wCmG5jtRc,1189
 brainstate/nn/_common.py,sha256=XQw0i0sH3Y_qUwHSMC7G9VQnDj-RuuTh1Ul-xRIPxxc,7136
-brainstate/nn/_exp_euler.py,sha256=cRgPNcjMs2C9x_8JabtYz5hm_FwqbiJ_U1VfRHYIlrE,3519
+brainstate/nn/_exp_euler.py,sha256=s-Z_cT_oYvCvE-OaXuUidIxQs3KOy1pzkx1lwtfPo00,3529
 brainstate/nn/_exp_euler_test.py,sha256=kvPf009DMYtla2uedKVKrPTHDyMTBepjlfsk5vDHqhI,1240
 brainstate/nn/_module.py,sha256=vrukVI0ylbymzilh9BZtb-d9dnsBsykqanUNTx9Eb6Y,12844
 brainstate/nn/_module_test.py,sha256=UrVA85fo0KVFN9ApPkxkRcvtXEskWOXPzZIBa4JSFo0,8891
@@ -66,7 +66,7 @@ brainstate/nn/metrics.py,sha256=p7eVwd5y8r0N5rMws-zOS_KaZCLOMdrXyQvLnoJeq1w,1473
 brainstate/nn/_dyn_impl/__init__.py,sha256=Oazar7h89dp1WA2Vx4Tj7gCBhxJKH4LAUEABkBEG7vU,1462
 brainstate/nn/_dyn_impl/_dynamics_neuron.py,sha256=mcDxVZlk56NAEkR6xcE74hOZ9up8Rua4SvKEeAhJKU4,10925
 brainstate/nn/_dyn_impl/_dynamics_neuron_test.py,sha256=_wPp6UvWVZI9EYba-DWL_JZXyMxm0-SHDkZHI8lKp8w,6315
-brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=Nk_c62nCJjjjhihQV_xqYeT-x_34xzc6OhdScO19ffw,15267
+brainstate/nn/_dyn_impl/_dynamics_synapse.py,sha256=xXMNlDWX0tQ9N0zORfT4DFKoEXtrsRbeetqOq-bYovs,15518
 brainstate/nn/_dyn_impl/_dynamics_synapse_test.py,sha256=VUDMlHNcyeqHrBd1eAXg_VD0HCSg5C-eqMmcJVzYcJA,4979
 brainstate/nn/_dyn_impl/_inputs.py,sha256=ubM5Z2q0gXpJ2M5Das3A5MJpFOorVomfW6-079mqJ9k,12935
 brainstate/nn/_dyn_impl/_projection_alignpost.py,sha256=PNC1Tzx_SF2DHAHeJCufXzO_Q4qLoBpWABI45B3GRuc,876
@@ -102,7 +102,7 @@ brainstate/optim/_lr_scheduler.py,sha256=Ua_H3VWUt9QZ0pHtGpnq2wrqsTOUidXJ3VDz9s-
 brainstate/optim/_lr_scheduler_test.py,sha256=W0F1eHb9S4seE468c26owcJIWTtgNhZYrOi2GrysVNI,1774
 brainstate/optim/_optax_optimizer.py,sha256=SuXV_xUBfhOw1_C2J5TIpy3dXDtI9VJFaSMLy8hLcXE,5312
 brainstate/optim/_optax_optimizer_test.py,sha256=J7zvmeSaWmTlfbpjx1ILb9cSC5qlj1wn4H2QMw9jUY0,1760
-brainstate/optim/_sgd_optimizer.py,sha256=oOFUEqCFX-WfhMnB614AjScbgFYP8y-zIwb94FZ_olA,46006
+brainstate/optim/_sgd_optimizer.py,sha256=RDQLrWsJFeEpTc93toDbf4hXnHmfY0hpTs38z1AZYPY,46144
 brainstate/random/__init__.py,sha256=c5q-RC3grRIjx-HBb2IhKZpi_xzbFmUUxzRAzqfREic,1045
 brainstate/random/_rand_funs.py,sha256=fmyaBTl_P7M63RRp5V2dOR_ttXQ0vB2qU2C8RqfkjY0,137680
 brainstate/random/_rand_funs_test.py,sha256=Nhy5gXuJ2ld9u8CTCqU1V94FPm0TvYQ-oMy2bP_CZ7I,19436
@@ -121,8 +121,8 @@ brainstate/util/_pretty_table.py,sha256=NM_6VAW6oL9jojsK0-RkQGHnDzLy_fn_hgzl5R8o
 brainstate/util/_scaling.py,sha256=pc_eM_SZVwkY65I4tJh1ODiHNCoEhsfFXl2zBK0PLAg,7562
 brainstate/util/_struct.py,sha256=F5GfFURITAIYTwf17_xypkZU1wvoL4dUCviPnr_eCtw,17515
 brainstate/util/filter.py,sha256=Zw0H42NwAi2P7dBr3ISv2VpkB5jqoWnV4Kpd61gq66o,14126
-brainstate-0.1.0.post20250218.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250218.dist-info/METADATA,sha256=INsdNphwnxGh07Urn2gjiOYzPEEQ8Lcr19bugHUYrU4,3585
-brainstate-0.1.0.post20250218.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250218.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250218.dist-info/RECORD,,
+brainstate-0.1.0.post20250315.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250315.dist-info/METADATA,sha256=8gOEdv6PiXBLr_gAvx70Yik7G7XxidMVPPLOLx3ndPc,3689
+brainstate-0.1.0.post20250315.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250315.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250315.dist-info/RECORD,,