brainstate 0.1.0.post20250222__py2.py3-none-any.whl → 0.1.0.post20250322__py2.py3-none-any.whl

brainstate/augment/_mapping.py

@@ -918,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,
             )
@@ -932,7 +932,7 @@ def map(
     return ys
 
 
-def flatten_(x):
+def _flatten(x):
     return x.reshape(-1, *x.shape[2:])
 
 
@@ -948,8 +948,13 @@ 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``?
+    if isinstance(axis_size, int) and axis_size <= 0:
+        raise ValueError(f"axis_size must be greater than 0, got {axis_size}.")
+
 
     @vmap(
         in_axes=in_axes,
@@ -957,6 +962,8 @@ def _vmap_new_states_transform(
         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
@@ -999,6 +1006,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.
@@ -1030,6 +1039,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(
@@ -1041,4 +1052,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,10 +15,11 @@
 
 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
@@ -268,6 +269,7 @@ class TestVmap(unittest.TestCase):
     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)
 
@@ -275,6 +277,42 @@ class TestVmap(unittest.TestCase):
         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):
@@ -360,3 +398,38 @@ class TestRemoveAxis:
         complex_array = jnp.array([[1 + 1j, 2 + 2j], [3 + 3j, 4 + 4j]])
         complex_result = _remove_axis(complex_array, 0)
         assert jnp.allclose(complex_result, jnp.array([1 + 1j, 2 + 2j]))
+
+
+class TestVMAPNewStatesEdgeCases(unittest.TestCase):
+
+    def test_axis_size_zero(self):
+        foo = brainstate.nn.LIF(3)
+        # Testing that axis_size of 0 raises an error.
+        with self.assertRaises(ValueError):
+            @bst.augment.vmap_new_states(state_tag='new1', axis_size=0)
+            def faulty_init():
+                foo.init_state()
+
+            # Call the decorated function to trigger validation
+            faulty_init()
+
+    def test_axis_size_negative(self):
+        foo = brainstate.nn.LIF(3)
+        # Testing that a negative axis_size raises an error.
+        with self.assertRaises(ValueError):
+            @bst.augment.vmap_new_states(state_tag='new1', axis_size=-3)
+            def faulty_init():
+                foo.init_state()
+
+            faulty_init()
+
+    def test_incompatible_shapes(self):
+        foo = brainstate.nn.LIF(3)
+        # Simulate an incompatible shapes scenario:
+        # We intentionally assign a state with a different shape than expected.
+        @bst.augment.vmap_new_states(state_tag='new1', axis_size=5)
+        def faulty_init():
+            # Modify state to produce an incompatible shape
+            foo.c = bst.State(jnp.arange(3))  # Original expected shape is (4,)
+
+        faulty_init()

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.
@@ -252,9 +250,9 @@ def vmap_call_all_functions(
 @set_module_as('brainstate.nn')
 def init_all_states(
     target: T,
-    init_args: Tuple[Any, ...] | Any = (),
-    init_kwargs: Dict[str, Any] | None = None,
+    *init_args,
     node_to_exclude: Filter = None,
+    **init_kwargs,
 ) -> T:
     """
     Initialize all states for the given target module and its submodules.
@@ -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/_collective_ops_test.py

@@ -34,3 +34,14 @@ class Test_vmap_init_all_states:
             print(gru)
 
         init()
+
+
+class Test_init_all_states:
+    def test_init_all_states(self):
+        gru = bst.nn.GRUCell(1, 2)
+        bst.nn.init_all_states(gru, batch_size=10)
+        print(gru)
+
+
+
+

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-0.1.0.post20250222.dist-info → brainstate-0.1.0.post20250322.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250222
+Version: 0.1.0.post20250322
 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.post20250222.dist-info → brainstate-0.1.0.post20250322.dist-info}/RECORD

@@ -14,8 +14,8 @@ brainstate/augment/_autograd.py,sha256=hfDoa2HbkRn-InOS0yOcb6gEZ2DLNqtWA133P8-hv
 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=X6C7IcYYXqDkYlWRXE_pdZIcCXvnltyNdESBoEm2fms,42940
-brainstate/augment/_mapping_test.py,sha256=M7pcXnRwBe2LyJh7VsWROGLR3N2bRfMUU1wj0ivbjNQ,12368
+brainstate/augment/_mapping.py,sha256=ru4byvHCtEvmSBTC-DckgKqyJEbKxstPdJcz3uE5AZ8,43494
+brainstate/augment/_mapping_test.py,sha256=Ax9-NjnCHPrvO_fPN24mOtl0NZugWg3AcNyxUYnHS1E,14709
 brainstate/augment/_random.py,sha256=ikRzNoDDE2BkARajDsBhNlngCUrghzGSZUDmEGvVors,5386
 brainstate/compile/__init__.py,sha256=fQtG316MLkeeu1Ssp54Kghw1PwbGK5gNq9yRVJu0wjA,1474
 brainstate/compile/_ad_checkpoint.py,sha256=3wv-f89oo94XeWwRV5LcRot0Nz7xTk5_PdjEDyUMsoo,9394
@@ -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_test.py,sha256=yW7NNYsGFglFRFkqVlpGSY6WLnU-h8GlK6wCmG5jtRc,1189
+brainstate/nn/_collective_ops.py,sha256=Hb8ruvqlhSA-XWJ66ReItyaIhxyWlSKGFZ9-EYMQ-mk,17457
+brainstate/nn/_collective_ops_test.py,sha256=nloqrlf6M7H-mgvHmIARrKzMotp8khxEuYSMPvXM5J0,1375
 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
@@ -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.post20250222.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
-brainstate-0.1.0.post20250222.dist-info/METADATA,sha256=Rzs6V20AVpOxQhtd8EJLidBkF7UJriVF4g2zkr0S92M,3585
-brainstate-0.1.0.post20250222.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
-brainstate-0.1.0.post20250222.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
-brainstate-0.1.0.post20250222.dist-info/RECORD,,
+brainstate-0.1.0.post20250322.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250322.dist-info/METADATA,sha256=ixPeEJ2tz1Dti9D7RFHS-bDLCj1a5VtVnbo7GS82xaE,3689
+brainstate-0.1.0.post20250322.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250322.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250322.dist-info/RECORD,,