brainstate 0.1.4__py2.py3-none-any.whl → 0.1.6__py2.py3-none-any.whl

brainstate/__init__.py

@@ -17,7 +17,7 @@
 A ``State``-based Transformation System for Program Compilation and Augmentation
 """
 
-__version__ = "0.1.4"
+__version__ = "0.1.6"
 
 from . import augment
 from . import compile

brainstate/_state.py

@@ -50,6 +50,7 @@ __all__ = [
     'LongTermState',
     'HiddenState',
     'ParamState',
+    'BatchState',
     'TreefyState',
     'FakeState',
 
@@ -115,10 +116,10 @@ def check_state_value_tree(val: bool = True) -> Generator[None, None, None]:
 
     Example::
 
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
-      >>> state = bst.ShortTermState(jnp.zeros((2, 3)))
-      >>> with bst.check_state_value_tree():
+      >>> state = brainstate.ShortTermState(jnp.zeros((2, 3)))
+      >>> with brainstate.check_state_value_tree():
       >>>   # The line below will not raise an error.
       >>>   state.value = jnp.zeros((2, 3))
       ...
@@ -162,10 +163,10 @@ def check_state_jax_tracer(val: bool = True) -> Generator[None, None, None]:
     Example::
 
       >>> import jax
-      >>> import brainstate as bst
+      >>> import brainstate as brainstate
       >>> import jax.numpy as jnp
       >>>
-      >>> a = bst.ShortTermState(jnp.zeros((2, 3)))
+      >>> a = brainstate.ShortTermState(jnp.zeros((2, 3)))
       >>>
       >>> @jax.jit
       >>> def run_state(b):

brainstate/augment/_autograd.py

@@ -173,6 +173,7 @@ class GradientTransform(PrettyRepr):
         return_value: bool = False,
         has_aux: bool = False,
         transform_params: Optional[Dict[str, Any]] = None,
+        check_states: bool = True,
     ):
         """
         Initialize a ``GradientTransform`` instance.
@@ -192,11 +193,12 @@ class GradientTransform(PrettyRepr):
         # gradient variables
         if isinstance(grad_states, dict):
             grad_states = {k: v for k, v in grad_states.items()}
-        self._grad_states, self._grad_tree = jax.tree.flatten(grad_states)
+        self._grad_states, self._grad_tree = jax.tree.flatten(grad_states, is_leaf=lambda x: isinstance(x, State))
         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.")
+        self.check_states = check_states
 
         # parameters
         if argnums is None and len(self._grad_states) == 0:
@@ -259,10 +261,15 @@ class GradientTransform(PrettyRepr):
             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)))
+            if self.check_states:
+                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)))
+            else:
+                id2state = {id(st): st for st in self._grad_states}
+                for id_ in all_ids:
+                    grad_vals[id_] = id2state[id_].value
 
         return grad_vals, other_vals
 
@@ -449,6 +456,7 @@ def grad(
     has_aux: Optional[bool] = None,
     return_value: Optional[bool] = False,
     unit_aware: bool = False,
+    check_states: bool = True,
 ) -> GradientTransform | Callable[[Callable], GradientTransform]:
     """
     Compute the gradient of a scalar-valued function with respect to its arguments.
@@ -493,7 +501,8 @@ def grad(
                 argnums=argnums,
                 return_value=return_value,
                 has_aux=False if has_aux is None else has_aux,
-                transform_params=dict(holomorphic=holomorphic, allow_int=allow_int)
+                transform_params=dict(holomorphic=holomorphic, allow_int=allow_int),
+                check_states=check_states
             )
 
         return transform
@@ -505,7 +514,8 @@ def grad(
         argnums=argnums,
         return_value=return_value,
         has_aux=False if has_aux is None else has_aux,
-        transform_params=dict(holomorphic=holomorphic, allow_int=allow_int)
+        transform_params=dict(holomorphic=holomorphic, allow_int=allow_int),
+        check_states=check_states
     )
 
 
@@ -520,6 +530,7 @@ def vector_grad(
     return_value: bool = False,
     has_aux: Optional[bool] = None,
     unit_aware: bool = False,
+    check_states: bool = True,
 ) -> GradientTransform | Callable[[Callable], GradientTransform]:
     """Take vector-valued gradients for function ``func``.
 
@@ -559,7 +570,8 @@ def vector_grad(
                 grad_states=grad_states,
                 argnums=argnums,
                 return_value=return_value,
-                has_aux=False if has_aux is None else has_aux
+                has_aux=False if has_aux is None else has_aux,
+                check_states=check_states
             )
 
         return transform
@@ -571,7 +583,8 @@ def vector_grad(
             grad_states=grad_states,
             argnums=argnums,
             return_value=return_value,
-            has_aux=False if has_aux is None else has_aux
+            has_aux=False if has_aux is None else has_aux,
+            check_states=check_states
         )
 
 
@@ -588,6 +601,7 @@ def jacrev(
     holomorphic: bool = False,
     allow_int: bool = False,
     unit_aware: bool = False,
+    check_states: bool = True,
 ) -> GradientTransform:
     """
     Extending automatic Jacobian (reverse-mode) of ``func`` to classes.
@@ -638,7 +652,8 @@ def jacrev(
         has_aux=False if has_aux is None else has_aux,
         transform_params=dict(holomorphic=holomorphic,
                               allow_int=allow_int,
-                              unit_aware=unit_aware, )
+                              unit_aware=unit_aware, ),
+        check_states=check_states
     )
 
 
@@ -656,6 +671,7 @@ def jacfwd(
     return_value: bool = False,
     holomorphic: bool = False,
     unit_aware: bool = False,
+    check_states: bool = True,
 ) -> GradientTransform:
     """Extending automatic Jacobian (forward-mode) of ``func`` to classes.
 
@@ -696,7 +712,8 @@ def jacfwd(
         argnums=argnums,
         return_value=return_value,
         has_aux=False if has_aux is None else has_aux,
-        transform_params=dict(holomorphic=holomorphic, unit_aware=unit_aware)
+        transform_params=dict(holomorphic=holomorphic, unit_aware=unit_aware),
+        check_states=check_states
     )
 
 
@@ -712,6 +729,7 @@ def hessian(
     holomorphic: bool = False,
     has_aux: Optional[bool] = None,
     unit_aware: bool = False,
+    check_states: bool = True,
 ) -> GradientTransform:
     """
     Hessian of ``func`` as a dense array.
@@ -752,7 +770,8 @@ def hessian(
         argnums=argnums,
         return_value=return_value,
         has_aux=False if has_aux is None else has_aux,
-        transform_params=dict(holomorphic=holomorphic)
+        transform_params=dict(holomorphic=holomorphic),
+        check_states=check_states
     )
 
 

brainstate/augment/_autograd_test.py

@@ -619,7 +619,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -635,12 +635,12 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacfwd(f1)(_x, _y)
 #     t = Test()
-#     br = bst.augment.jacfwd(t, grad_states=t.x)()
+#     br = brainstate.augment.jacfwd(t, grad_states=t.x)()
 #     self.assertTrue((br == _jr).all())
 #
 #     _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
 #     t = Test()
-#     br = bst.augment.jacfwd(t, grad_states=[t.x, t.y])()
+#     br = brainstate.augment.jacfwd(t, grad_states=[t.x, t.y])()
 #     self.assertTrue((br[0] == _jr[0]).all())
 #     self.assertTrue((br[1] == _jr[1]).all())
 #
@@ -652,7 +652,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -667,12 +667,12 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacrev(f1)(_x, _y)
 #     t = Test()
-#     br = bst.augment.jacrev(t, grad_states=t.x)(_y)
+#     br = brainstate.augment.jacrev(t, grad_states=t.x)(_y)
 #     self.assertTrue((br == _jr).all())
 #
 #     _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
 #     t = Test()
-#     var_grads, arg_grads = bst.augment.jacrev(t, grad_states=t.x, argnums=0)(_y)
+#     var_grads, arg_grads = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0)(_y)
 #     print(var_grads, )
 #     print(arg_grads, )
 #     self.assertTrue((var_grads == _jr[0]).all())
@@ -686,7 +686,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -701,12 +701,12 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacfwd(f1)(_x, _y)
 #     t = Test()
-#     br = bst.augment.jacfwd(t, grad_states=t.x)(_y)
+#     br = brainstate.augment.jacfwd(t, grad_states=t.x)(_y)
 #     self.assertTrue((br == _jr).all())
 #
 #     _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
 #     t = Test()
-#     var_grads, arg_grads = bst.augment.jacfwd(t, grad_states=t.x, argnums=0)(_y)
+#     var_grads, arg_grads = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0)(_y)
 #     print(var_grads, )
 #     print(arg_grads, )
 #     self.assertTrue((var_grads == _jr[0]).all())
@@ -722,7 +722,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -737,13 +737,13 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacrev(f1)(_x, _y)
 #     t = Test()
-#     br, _ = bst.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
+#     br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
 #     self.assertTrue((br == _jr).all())
 #
 #     t = Test()
 #     _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
 #     _aux = t(_y)[1]
-#     (var_grads, arg_grads), aux = bst.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
+#     (var_grads, arg_grads), aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
 #     print(var_grads, )
 #     print(arg_grads, )
 #     self.assertTrue((var_grads == _jr[0]).all())
@@ -762,7 +762,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -777,7 +777,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacfwd(f1)(_x, _y)
 #     t = Test()
-#     br, (c, d) = bst.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
+#     br, (c, d) = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
 #     # print(_jr)
 #     # print(br)
 #     a = (br == _jr)
@@ -786,7 +786,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     t = Test()
 #     _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
 #     _aux = t(_y)[1]
-#     (var_grads, arg_grads), aux = bst.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
+#     (var_grads, arg_grads), aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True)(_y)
 #     print(var_grads, )
 #     print(arg_grads, )
 #     self.assertTrue((var_grads == _jr[0]).all())
@@ -805,7 +805,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -820,13 +820,13 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacrev(f1)(_x, _y)
 #     t = Test()
-#     br, _ = bst.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
+#     br, _ = brainstate.augment.jacrev(t, grad_states=t.x, has_aux=True)(_y)
 #     self.assertTrue((br == _jr).all())
 #
 #     t = Test()
 #     _jr = jax.jacrev(f1, argnums=(0, 1))(_x, _y)
 #     _val, _aux = t(_y)
-#     (var_grads, arg_grads), value, aux = bst.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
+#     (var_grads, arg_grads), value, aux = brainstate.augment.jacrev(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
 #     print(var_grads, )
 #     print(arg_grads, )
 #     self.assertTrue((var_grads == _jr[0]).all())
@@ -846,7 +846,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.array([1., 2., 3.])
 #     _y = jnp.array([10., 5.])
 #
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.array([1., 2., 3.]))
@@ -861,13 +861,13 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     _jr = jax.jacfwd(f1)(_x, _y)
 #     t = Test()
-#     br, _ = bst.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
+#     br, _ = brainstate.augment.jacfwd(t, grad_states=t.x, has_aux=True)(_y)
 #     self.assertTrue((br == _jr).all())
 #
 #     t = Test()
 #     _jr = jax.jacfwd(f1, argnums=(0, 1))(_x, _y)
 #     _val, _aux = t(_y)
-#     (var_grads, arg_grads), value, aux = bst.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
+#     (var_grads, arg_grads), value, aux = brainstate.augment.jacfwd(t, grad_states=t.x, argnums=0, has_aux=True, return_value=True)(_y)
 #     print(_val, )
 #     print('_aux: ', _aux, 'aux: ', aux)
 #     print(var_grads, )
@@ -884,7 +884,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #   def test1(self):
 #     f = lambda x: 3 * x ** 2
 #     _x = jnp.ones(10)
-#     pprint(bst.augment.vector_grad(f, argnums=0)(_x))
+#     pprint(brainstate.augment.vector_grad(f, argnums=0)(_x))
 #
 #   def test2(self):
 #     def f(x, y):
@@ -894,14 +894,14 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones(5)
 #     _y = jnp.ones(5)
 #
-#     g = bst.augment.vector_grad(f, argnums=0)(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
 #     pprint(g)
 #     self.assertTrue(jnp.array_equal(g, 2 * _x))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0,))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
 #     pprint(g)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x))
 #     self.assertTrue(jnp.array_equal(g[1], 2 * _y))
@@ -915,14 +915,14 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones(5)
 #     _y = jnp.ones(5)
 #
-#     g = bst.augment.vector_grad(f, argnums=0)(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
 #     # pprint(g)
 #     self.assertTrue(jnp.array_equal(g, 2 * _x + 3 * _x ** 2))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0,))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
 #     # pprint(g)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x + 3 * _x ** 2))
 #     self.assertTrue(jnp.array_equal(g[1], 2 * _y + 3 * _y ** 2))
@@ -935,14 +935,14 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones((5, 5))
 #     _y = jnp.ones((5, 5))
 #
-#     g = bst.augment.vector_grad(f, argnums=0)(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=0)(_x, _y)
 #     pprint(g)
 #     self.assertTrue(jnp.array_equal(g, 2 * _x))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0,))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0,))(_x, _y)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x))
 #
-#     g = bst.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
+#     g = brainstate.augment.vector_grad(f, argnums=(0, 1))(_x, _y)
 #     pprint(g)
 #     self.assertTrue(jnp.array_equal(g[0], 2 * _x))
 #     self.assertTrue(jnp.array_equal(g[1], 2 * _y))
@@ -956,7 +956,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones(5)
 #     _y = jnp.ones(5)
 #
-#     g, aux = bst.augment.vector_grad(f, has_aux=True)(_x, _y)
+#     g, aux = brainstate.augment.vector_grad(f, has_aux=True)(_x, _y)
 #     pprint(g, )
 #     pprint(aux)
 #     self.assertTrue(jnp.array_equal(g, 2 * _x))
@@ -970,7 +970,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones(5)
 #     _y = jnp.ones(5)
 #
-#     g, value = bst.augment.vector_grad(f, return_value=True)(_x, _y)
+#     g, value = brainstate.augment.vector_grad(f, return_value=True)(_x, _y)
 #     pprint(g, )
 #     pprint(value)
 #     self.assertTrue(jnp.array_equal(g, 2 * _x))
@@ -985,7 +985,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #     _x = jnp.ones(5)
 #     _y = jnp.ones(5)
 #
-#     g, value, aux = bst.augment.vector_grad(f, has_aux=True, return_value=True)(_x, _y)
+#     g, value, aux = brainstate.augment.vector_grad(f, has_aux=True, return_value=True)(_x, _y)
 #     print('grad', g)
 #     print('value', value)
 #     print('aux', aux)
@@ -996,7 +996,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 # class TestClassFuncVectorGrad(unittest.TestCase):
 #   def test1(self):
-#     class Test(bst.nn.Module):
+#     class Test(brainstate.nn.Module):
 #       def __init__(self):
 #         super(Test, self).__init__()
 #         self.x = jnp.Variable(jnp.ones(5))
@@ -1007,13 +1007,13 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #     t = Test()
 #
-#     g = bst.augment.vector_grad(t, grad_states=t.x)()
+#     g = brainstate.augment.vector_grad(t, grad_states=t.x)()
 #     self.assertTrue(jnp.array_equal(g, 2 * t.x))
 #
-#     g = bst.augment.vector_grad(t, grad_states=(t.x,))()
+#     g = brainstate.augment.vector_grad(t, grad_states=(t.x,))()
 #     self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
 #
-#     g = bst.augment.vector_grad(t, grad_states=(t.x, t.y))()
+#     g = brainstate.augment.vector_grad(t, grad_states=(t.x, t.y))()
 #     self.assertTrue(jnp.array_equal(g[0], 2 * t.x))
 #     self.assertTrue(jnp.array_equal(g[1], 2 * t.y))
 #
@@ -1025,20 +1025,20 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 # class TestDebug(parameterized.TestCase):
 #   def test_debug1(self):
-#     a = bst.random.RandomState()
+#     a = brainstate.random.RandomState()
 #
 #     def f(b):
 #       print(a.value)
 #       return a + b + a.random()
 #
-#     f = bst.augment.vector_grad(f, argnums=0)
+#     f = brainstate.augment.vector_grad(f, argnums=0)
 #     f(1.)
 #
 #     with jax.disable_jit():
 #       f(1.)
 #
 #   @parameterized.product(
-#     grad_fun=[bst.augment.grad, bst.augment.vector_grad]
+#     grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
 #   )
 #   def test_print_info1(self, grad_fun):
 #     file = tempfile.TemporaryFile(mode='w+')
@@ -1075,7 +1075,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #       self.assertTrue(file.read().strip() == expect_res.strip())
 #
 #   @parameterized.product(
-#     grad_fun=[bst.augment.grad, bst.augment.vector_grad]
+#     grad_fun=[brainstate.augment.grad, brainstate.augment.vector_grad]
 #   )
 #   def test_print_info2(self, grad_fun):
 #     file = tempfile.TemporaryFile(mode='w+')
@@ -1117,7 +1117,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #       a = jnp.Variable(jnp.ones(2))
 #       b = jnp.Variable(jnp.zeros(2))
 #
-#       @bst.augment.vector_grad(argnums=0)
+#       @brainstate.augment.vector_grad(argnums=0)
 #       def f1(c):
 #         a.value += 1
 #         b.value += 10
@@ -1149,7 +1149,7 @@ class TestClassFuncJacobian(unittest.TestCase):
 #
 #
 #     def run_fun(d):
-#       @bst.augment.vector_grad(argnums=0)
+#       @brainstate.augment.vector_grad(argnums=0)
 #       def f1(c):
 #         a.value += d
 #         b.value += 10
@@ -1178,8 +1178,8 @@ class TestClassFuncJacobian(unittest.TestCase):
 #         print('compiling f ...', file=file)
 #         return a + b
 #
-#       grad1 = bst.augment.grad(f)(1., 2.)  # call "f" twice, one for Variable finding, one for compiling
-#       grad2 = bst.augment.vector_grad(f)(1., 2.)  # call "f" once for compiling
+#       grad1 = brainstate.augment.grad(f)(1., 2.)  # call "f" twice, one for Variable finding, one for compiling
+#       grad2 = brainstate.augment.vector_grad(f)(1., 2.)  # call "f" once for compiling
 #
 #       file.seek(0)
 #       print(file.read().strip())

brainstate/augment/_eval_shape.py

@@ -40,13 +40,13 @@ def abstract_init(
 
     Here's an example::
 
-        >>> import brainstate as bst
+        >>> import brainstate
         >>> class MLP:
         ...     def __init__(self, n_in, n_mid, n_out):
-        ...         self.dense1 = bst.nn.Linear(n_in, n_mid)
-        ...         self.dense2 = bst.nn.Linear(n_mid, n_out)
+        ...         self.dense1 = brainstate.nn.Linear(n_in, n_mid)
+        ...         self.dense2 = brainstate.nn.Linear(n_mid, n_out)
 
-        >>> r = bst.augment.abstract_init(lambda: MLP(1, 2, 3))
+        >>> r = brainstate.augment.abstract_init(lambda: MLP(1, 2, 3))
         >>> r
         MLP(
           dense1=Linear(

brainstate/augment/_mapping.py

@@ -185,10 +185,10 @@ def _compile_stateful_function(
     if isinstance(in_axes, int):
         args = jax.tree.map(lambda x: _remove_axis(x, in_axes), args)
     elif isinstance(in_axes, tuple):
-        args = tuple(
-            [arg if in_axis is None else _remove_axis(arg, in_axis)
-             for arg, in_axis in zip(args, in_axes)]
-        )
+        args = tuple([
+            arg if in_axis is None else _remove_axis(arg, in_axis)
+            for arg, in_axis in zip(args, in_axes)
+        ])
     stateful_fn.make_jaxpr(state_vals, args)
     return stateful_fn.get_arg_cache_key(state_vals, args)
 
@@ -383,10 +383,7 @@ def _vmap_transform(
         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(
@@ -460,7 +457,10 @@ def _vmap_transform(
         # analyze vmapping axis error
         for state in state_trace.get_write_states():
             leaves = jax.tree.leaves(state.value)
-            if any([isinstance(leaf, BatchTracer) for leaf in leaves]) and state not in out_state_to_axis:
+            if (
+                any([isinstance(leaf, BatchTracer) and (leaf.batch_dim is not None) for leaf in leaves])
+                and state not in out_state_to_axis
+            ):
                 if isinstance(state, RandomState) and state in rng_sets:
                     continue
                 state.raise_error_with_source_info(
@@ -481,7 +481,7 @@ def _vmap_transform(
         return out_rng_keys, out_state_vmap_vals, out_state_oth_vals, outs
 
     @functools.wraps(f)
-    def vmapped_fn(*args):
+    def vmapped_fn(*args, **kwargs):
         """
         Applies vectorized mapping (vmap) to the input function while managing state.
 
@@ -503,6 +503,11 @@ def _vmap_transform(
             data structures (e.g., axis_to_in_states, in_state_to_axis) which
             should be defined in the broader context.
         """
+        if len(kwargs):
+            raise NotImplementedError(
+                "Keyword arguments `f(**kwargs)` are not supported in brainstate.augment.vmap"
+            )
+
         # in states values
         in_state_map_vals = [
             [st.value for st in states]
@@ -615,16 +620,16 @@ def vmap(
 
     These are several example usage::
 
-        >>> import brainstate as bst
+        >>> import brainstate as brainstate
         >>> import jax.numpy as jnp
 
-        >>> class Model(bst.nn.Module):
+        >>> class Model(brainstate.nn.Module):
         >>>     def __init__(self):
         >>>         super().__init__()
         >>>
-        >>>         self.a = bst.ShortTermState(bst.random.randn(5))
-        >>>         self.b = bst.ShortTermState(bst.random.randn(5))
-        >>>         self.c = bst.State(bst.random.randn(1))
+        >>>         self.a = brainstate.ShortTermState(brainstate.random.randn(5))
+        >>>         self.b = brainstate.ShortTermState(brainstate.random.randn(5))
+        >>>         self.c = brainstate.State(brainstate.random.randn(1))
 
         >>>     def __call__(self, *args, **kwargs):
         >>>         self.c.value = self.a.value * self.b.value
@@ -632,9 +637,9 @@ def vmap(
 
         >>> model = Model()
 
-        >>> r = bst.augment.vmap(
+        >>> r = brainstate.augment.vmap(
         >>>     model,
-        >>>     in_states=model.states(bst.ShortTermState),
+        >>>     in_states=model.states(brainstate.ShortTermState),
         >>>     out_states=model.c
         >>> )()