brainstate 0.1.0.post20250218__py2.py3-none-any.whl → 0.1.0.post20250222__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:
@@ -936,10 +949,8 @@ def _vmap_new_states_transform(
     state_tag: str | None = None,
     state_to_exclude: Filter | None = None,
 ):
-
     # TODO: How about nested call ``vmap_new_states``?
 
-
     @vmap(
         in_axes=in_axes,
         out_axes=out_axes,

brainstate/augment/_mapping_test.py

@@ -21,6 +21,7 @@ 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 +265,17 @@ 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))
+
+
 
 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/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.post20250222.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: brainstate
-Version: 0.1.0.post20250218
+Version: 0.1.0.post20250222
 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.post20250218.dist-info → brainstate-0.1.0.post20250222.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=X6C7IcYYXqDkYlWRXE_pdZIcCXvnltyNdESBoEm2fms,42940
+brainstate/augment/_mapping_test.py,sha256=M7pcXnRwBe2LyJh7VsWROGLR3N2bRfMUU1wj0ivbjNQ,12368
 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
@@ -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.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,,