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

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

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

@@ -54,7 +54,7 @@ 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=Hb8ruvqlhSA-XWJ66ReItyaIhxyWlSKGFZ9-EYMQ-mk,17457
+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
@@ -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.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,,
+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,,