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

brainstate/augment/_mapping.py

@@ -952,6 +952,9 @@ def _vmap_new_states_transform(
     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,

brainstate/augment/_mapping_test.py

@@ -315,7 +315,6 @@ class TestVmap(unittest.TestCase):
         )
 
 
-
 class TestMap(unittest.TestCase):
     def test_map(self):
         for dim in [(10,), (10, 10), (10, 10, 10)]:
@@ -399,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

@@ -16,9 +16,9 @@
 from __future__ import annotations
 
 from collections import namedtuple
+from typing import Callable, TypeVar, Tuple, Any, Dict
 
 import jax
-from typing import Callable, TypeVar, Tuple, Any, Dict
 
 from brainstate._state import catch_new_states
 from brainstate._utils import set_module_as
@@ -250,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.
@@ -289,12 +289,12 @@ def init_all_states(
 @set_module_as('brainstate.nn')
 def vmap_init_all_states(
     target: T,
-    init_args: Tuple[Any, ...] | Any = (),
-    init_kwargs: Dict[str, Any] | None = None,
+    *init_args: Tuple[Any, ...] | Any,
     axis_size: int = None,
     node_to_exclude: Filter = None,
     state_to_exclude: Filter = None,
     state_tag: str | None = None,
+    **init_kwargs: Dict[str, Any] | None
 ) -> T:
     """
     Initialize all vmap states for the given target module.
@@ -342,8 +342,8 @@ def vmap_init_all_states(
     def init_fn():
         init_all_states(
             target,
-            init_args=init_args,
-            init_kwargs=init_kwargs,
+            *init_args,
+            **init_kwargs,
             node_to_exclude=node_to_exclude,
         )
         return

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

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

{brainstate-0.1.0.post20250315.dist-info → brainstate-0.1.0.post20250325.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=BPwpD7jX4xRNl4BdAsKGoF45MKbmEF9Lyyp11pJucIg,43356
-brainstate/augment/_mapping_test.py,sha256=-4HJXmJw_6SD9dQnHTBjgYVuq6VTVjz0xpc9v2CJVNw,13414
+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,8 +54,8 @@ 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=NI9BT-908TbIlXLMjbWsPyI5YLZD_cCkSKGeOY-qO60,17512
-brainstate/nn/_collective_ops_test.py,sha256=yW7NNYsGFglFRFkqVlpGSY6WLnU-h8GlK6wCmG5jtRc,1189
+brainstate/nn/_collective_ops.py,sha256=v5deEfjCWylCk4bV0b3mLHjFUJ0L9YiGhJ7D_RuZwBE,17428
+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=s-Z_cT_oYvCvE-OaXuUidIxQs3KOy1pzkx1lwtfPo00,3529
 brainstate/nn/_exp_euler_test.py,sha256=kvPf009DMYtla2uedKVKrPTHDyMTBepjlfsk5vDHqhI,1240
@@ -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.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,,
+brainstate-0.1.0.post20250325.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250325.dist-info/METADATA,sha256=Ouz5zN2P7bIG1OYkQRFOaSxWeL4OOjHRJ2cQKaqZc6E,3689
+brainstate-0.1.0.post20250325.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250325.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250325.dist-info/RECORD,,