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

brainstate/nn/_interaction/_normalizations.py

@@ -230,7 +230,7 @@ def _normalize(
         y = y * mul
         if weights is not None:
             y = weights.execute(y)
-        dtype = canonicalize_dtype(x, *jax.tree.leaves(weights.value), dtype=dtype)
+            dtype = canonicalize_dtype(x, *jax.tree.leaves(weights.value), dtype=dtype)
     else:
         assert var is None, 'mean and val must be both None or not None.'
         assert weights is None, 'scale and bias are not supported without mean and val'

brainstate/transform.py

@@ -15,3 +15,9 @@
 
 # alias for compilation and augmentation functions
 
+from .augment import *
+from .compile import *
+
+if __name__ == '__main__':
+    ifelse
+    grad

{brainstate-0.1.0.post20250325.dist-info → brainstate-0.1.0.post20250413.dist-info}/METADATA

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

@@ -7,7 +7,7 @@ brainstate/environ_test.py,sha256=khZ_-SUJL6rQCgndeYV98ruUIHGTwFDtITrOs_olmuo,20
 brainstate/mixin.py,sha256=g7uVUwZphZWsNs9pb48ozG2cDGaj0hs0g3lq8tDk-Sg,11310
 brainstate/mixin_test.py,sha256=Oq_0fwC9vpXDN4t4dTBhWzLdFDNlcYsrcip14F1yECI,3079
 brainstate/surrogate.py,sha256=wWYw-TxaFxHVneXuHjWD1UtTcOTk3XRSnhRtUkt_Hb8,53580
-brainstate/transform.py,sha256=vZWzO4F7qTsXL_SiVQPlTz0l9b_hRo9D-igETfgCTy0,758
+brainstate/transform.py,sha256=OvshYpPnp3YXPG6riy15Ve7jcX8t0aaATl1xZnsFeic,858
 brainstate/typing.py,sha256=988gX1tvwtyYnYjmej90OaRxoMoBIPO0-DSrXXGxojM,10523
 brainstate/augment/__init__.py,sha256=Q9-JIwQ1FNn8VLS1MA9MrSylbvUjWSw98whrI3NIuKo,1229
 brainstate/augment/_autograd.py,sha256=hfDoa2HbkRn-InOS0yOcb6gEZ2DLNqtWA133P8-hvIo,30138
@@ -92,7 +92,7 @@ brainstate/nn/_interaction/_conv_test.py,sha256=VJacu1afSmxFpgKerjMA2r3kwXAhf7aV
 brainstate/nn/_interaction/_embedding.py,sha256=iK0I1ExKWFa_QzV9UDGj32Ljsmdr1g_LlAtMcusebxU,2187
 brainstate/nn/_interaction/_linear.py,sha256=_CBAVR2EYE3tW60_nPq8pf7byuPoRsnqklC8iTF2hNU,14765
 brainstate/nn/_interaction/_linear_test.py,sha256=vQiXZDS2SMsbbQ45YFC9NHsVrv6ahwm-NoHZ1wenEF8,2870
-brainstate/nn/_interaction/_normalizations.py,sha256=g42vRntBC_vBmjhznXUksOmKVAXhaujHEQWwAxjOww4,37432
+brainstate/nn/_interaction/_normalizations.py,sha256=-wU1AAmz7p6hYmc1UIbm3MWieIkHMiXOXRJm5raiEbg,37436
 brainstate/nn/_interaction/_normalizations_test.py,sha256=2p1Jf8nA999VYGWbvOZfKYlKk6UmL0vaEB76xkXxkXw,2438
 brainstate/nn/_interaction/_poolings.py,sha256=Fk35jqUg09_4IumWZ8XFfku3-z5Vkqu0LT8ALeF4qYo,47018
 brainstate/nn/_interaction/_poolings_test.py,sha256=wmd5PngZ3E9tNyF3s0xk-DoDR5yFqpTi9A6nbNoIqn4,7429
@@ -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.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,,
+brainstate-0.1.0.post20250413.dist-info/LICENSE,sha256=VZe9u1jgUL2eCY6ZPOYgdb8KCblCHt8ECdbtJid6e1s,11550
+brainstate-0.1.0.post20250413.dist-info/METADATA,sha256=lCdOfnus0uYdv5QLOxASN8Mtkn70CkYY6YsQnAdBC-o,3689
+brainstate-0.1.0.post20250413.dist-info/WHEEL,sha256=bb2Ot9scclHKMOLDEHY6B2sicWOgugjFKaJsT7vwMQo,110
+brainstate-0.1.0.post20250413.dist-info/top_level.txt,sha256=eQbGgKn0ptx7FDWuua0V0wr4K1VHi2iOUCYo3fUQBRA,11
+brainstate-0.1.0.post20250413.dist-info/RECORD,,