brainstate 0.1.0.post20250105__py2.py3-none-any.whl → 0.1.0.post20250120__py2.py3-none-any.whl

brainstate/augment/_mapping_test.py

@@ -17,57 +17,122 @@ from __future__ import annotations
 
 import unittest
 
-import jax.core
 import jax.numpy as jnp
 
 import brainstate as bst
+from brainstate.augment._mapping import BatchAxisError
 
 
 class TestVmap(unittest.TestCase):
-    def test_vmap_return_keep_reference_return(self):
-        @bst.augment.vmap(in_axes=0, out_axes=0)
-        def create_model(key):
-            bst.random.set_key(key)
-            m1 = bst.nn.Linear(2, 3)
-
-            m2 = bst.nn.Linear(3, 4)
-            m2.a = m1
-            m3 = bst.nn.Linear(3, 5)
-            m3.a = m1
-            self.assertTrue(id(m2.a) == id(m3.a))
-            return m2, m3
-
-        m2, m3 = create_model(bst.random.split_key(10))
-        self.assertTrue(id(m2.a) == id(m3.a))
-        jax.core.concrete_or_error(None, bst.random.DEFAULT.value)
-
-    def test_vmap_return_keep_reference_pass_into_fun(self):
-        @bst.augment.vmap(in_axes=(None, None, 0), out_axes=0)
-        def run_model(m2, m3, x):
-            self.assertTrue(id(m2.a) == id(m3.a))
-            self.assertTrue(id(m2) != m2_id)
-            self.assertTrue(id(m3) != m3_id)
-            return m2(x), m3(x)
-
-        m1 = bst.nn.Linear(2, 3)
-        m2 = bst.nn.Linear(4, 3)
-        m2.a = m1
-        m3 = bst.nn.Linear(4, 5)
-        m3.a = m1
-        m3_id = id(m3)
-        m2_id = id(m2)
-        r1, r2 = run_model(m2, m3, jnp.ones((4, 3, 4)))
-
-    def test_vmap_set_key(self):
-        @bst.augment.vmap(in_axes=0, out_axes=0)
-        def create_model(key):
-            bst.random.set_key(key)
-            return bst.nn.Linear(2, 3)
-
-        model = create_model(bst.random.split_keys(10))
-        print(model.weight.value_call(jnp.shape))
-        model.weight.value_call(lambda x: jax.core.concrete_or_error(None, x))
-        bst.random.seed()
+    def test_vmap_1(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+                self.a = bst.State(bst.random.randn(5))
+                self.b = bst.State(bst.random.randn(5))
+
+            def __call__(self, *args, **kwargs):
+                return self.a.value * self.b.value
+
+        model = Model()
+        r1 = model.a.value * model.b.value
+        r2 = bst.augment.vmap(model, in_states=model.states())()
+        self.assertTrue(jnp.allclose(r1, r2))
+
+    def test_vmap_2(self):
+        class Model(bst.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))
+
+            def __call__(self, *args, **kwargs):
+                self.c.value = self.a.value * self.b.value
+                return self.c.value + 1.
+
+        model = Model()
+        with self.assertRaises(BatchAxisError):
+            r2 = bst.augment.vmap(model, in_states=model.states(bst.ShortTermState))()
+
+        model = Model()
+        r2 = bst.augment.vmap(model, in_states=model.states(bst.ShortTermState), out_states=model.c)()
+
+    def test_vmap_3(self):
+        class Model(bst.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+                self.a = bst.State(bst.random.randn(5))
+                self.b = bst.State(bst.random.randn(5))
+
+            def __call__(self, *args, **kwargs):
+                return self.a.value * self.b.value
+
+        model = Model()
+        with self.assertRaises(BatchAxisError):
+            r2 = bst.augment.vmap(model, in_states=model.states(), out_states={1: model.states()})()
+
+    def test_vmap_with_random(self):
+        class Model(bst.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))
+
+            def __call__(self, key):
+                bst.random.set_key(key)
+                self.c.value = self.a.value * self.b.value
+                return self.c.value + bst.random.randn(1)
+
+        model = Model()
+        r2 = bst.augment.vmap(
+            model,
+            in_states=model.states(bst.ShortTermState),
+            out_states=model.c
+        )(
+            bst.random.split_key(5)
+        )
+        print(bst.random.DEFAULT)
+
+    def test_vmap_with_random_2(self):
+        class Model(bst.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.rng = bst.random.RandomState(1)
+
+            def __call__(self, key):
+                self.rng.set_key(key)
+                self.c.value = self.a.value * self.b.value
+                return self.c.value + bst.random.randn(1)
+
+        model = Model()
+        with self.assertRaises(BatchAxisError):
+            r2 = bst.augment.vmap(
+                model,
+                in_states=model.states(bst.ShortTermState),
+                out_states=model.c
+            )(
+                bst.random.split_key(5)
+            )
+
+        model = Model()
+        r2 = bst.augment.vmap(
+            model,
+            in_states=model.states(bst.ShortTermState),
+            out_states=model.c,
+            rngs=model.rng,
+        )(
+            bst.random.split_key(5)
+        )
 
     def test_vmap_input(self):
         model = bst.nn.Linear(2, 3)
@@ -98,8 +163,8 @@ class TestVmap(unittest.TestCase):
 
         @bst.augment.vmap(in_axes=(None, 0), out_axes=0)
         def forward(model, x):
-            self.assertTrue(id(model) != model_id)
-            self.assertTrue(id(model.weight) != weight_id)
+            self.assertTrue(id(model) == model_id)
+            self.assertTrue(id(model.weight) == weight_id)
             print(id(model), id(model.weight))
             return model(x)
 
@@ -108,51 +173,7 @@ class TestVmap(unittest.TestCase):
         print(model.weight.value_call(jnp.shape))
         print(model.weight.value)
 
-    def test_vmap1(self):
-        model = bst.nn.Linear(2, 3)
-        x = jnp.ones((5, 2))
-
-        @bst.augment.vmap(in_axes=(None, 0), out_axes=0)
-        def forward(model, x):
-            return model(x)
-
-        y = forward(model, x)
-        print(y.shape)
-
-    def test_vmap2(self):
-        class LinearEnsemble(bst.nn.Module):
-            def __init__(self, num):
-                super().__init__()
-                self.w = bst.ParamState(bst.random.random((num, 2, 3)))
-
-        model = LinearEnsemble(5)
-        x = jnp.ones((2,))
-
-        @bst.augment.vmap(in_axes=(0, None), out_axes=0)
-        def forward(model, x):
-            return jnp.dot(x, model.w.value)
-
-        y = forward(model, x)
-        print(y.shape)
-
-    def test_vmap3(self):
-        class Foo(bst.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.a = bst.ParamState(jnp.arange(4))
-                self.b = bst.ShortTermState(jnp.arange(4))
-
-        state_axes = bst.augment.StateAxes({bst.ParamState: 0, bst.ShortTermState: None})
-
-        @bst.augment.vmap(in_axes=(state_axes,), out_axes=0)
-        def mul(foo):
-            return foo.a.value * foo.b.value
-
-        foo = Foo()
-        y = mul(foo)
-        print(y.shape)
-
-    def test_vmap4(self):
+    def test_vmap_jit(self):
         class Foo(bst.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -162,60 +183,35 @@ class TestVmap(unittest.TestCase):
             def __call__(self):
                 self.b.value = self.a.value * self.b.value
 
-        @bst.augment.vmap
-        def mul(foo):
-            foo()
-            return foo
-
         foo = Foo()
-        with bst.StateTraceStack() as trace:
-            m = mul(foo)
-
-        self.assertTrue(m is foo)
-        print(m.a.value, foo.a.value)
-        self.assertTrue(jnp.allclose(m.a.value, foo.a.value))
-        print(m.b.value, foo.b.value)
-        self.assertTrue(jnp.allclose(m.b.value, foo.b.value))
-        print(trace.get_write_states())
-        self.assertTrue(len(trace.get_write_states()) == 1)
-        print(trace.get_read_states())
-        self.assertTrue(len(trace.get_read_states()) == 2)
-
-    def test_vmap5(self):
-        class Foo(bst.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.a = bst.ParamState(jnp.arange(4))
-                self.b = bst.ShortTermState(jnp.arange(4))
 
-            def __call__(self):
-                self.b.value = self.a.value * self.b.value
-
-        @bst.augment.vmap
-        def mul(foo):
+        @bst.augment.vmap(in_states=foo.states())
+        def mul():
             foo()
 
-        foo = Foo()
+        @bst.compile.jit
+        def mul_jit(inp):
+            mul()
+            foo.a.value += inp
+
         with bst.StateTraceStack() as trace:
-            mul(foo)
+            mul_jit(1.)
 
         print(foo.a.value)
         print(foo.b.value)
-        self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4)))
+        self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4) + 1.))
         self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
 
         write_state_ids = [id(st) for st in trace.get_write_states()]
         read_state_ids = [id(st) for st in trace.get_read_states()]
 
-        assert id(foo.a) in read_state_ids
+        assert id(foo.a) in write_state_ids
         assert id(foo.b) in write_state_ids
 
         print(trace.get_write_states())
         print(trace.get_read_states())
 
-
-
-    def test_vmap_jit(self):
+    def test_vmap_jit_2(self):
         class Foo(bst.nn.Module):
             def __init__(self):
                 super().__init__()
@@ -225,70 +221,46 @@ class TestVmap(unittest.TestCase):
             def __call__(self):
                 self.b.value = self.a.value * self.b.value
 
-        @bst.augment.vmap
-        def mul(foo):
+        foo = Foo()
+
+        @bst.augment.vmap(in_states=foo.states())
+        def mul():
             foo()
 
         @bst.compile.jit
         def mul_jit(inp):
-            mul(foo)
-            foo.a.value += inp
+            mul()
+            foo.b.value += inp
 
-        foo = Foo()
         with bst.StateTraceStack() as trace:
             mul_jit(1.)
 
         print(foo.a.value)
         print(foo.b.value)
-        self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4) + 1.))
-        self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
+        self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4)))
+        self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4) + 1.))
 
         write_state_ids = [id(st) for st in trace.get_write_states()]
         read_state_ids = [id(st) for st in trace.get_read_states()]
 
-        assert id(foo.a) in write_state_ids
+        assert id(foo.a) in read_state_ids
         assert id(foo.b) in write_state_ids
 
         print(trace.get_write_states())
         print(trace.get_read_states())
 
 
-    def test_vmap_grad(self):
-        class Foo(bst.nn.Module):
-            def __init__(self):
-                super().__init__()
-                self.a = bst.ParamState(jnp.arange(4.))
-                self.b = bst.ShortTermState(jnp.arange(4.))
-
-            def __call__(self):
-                self.b.value = self.a.value * self.b.value
-
-        @bst.augment.vmap
-        def mul(foo):
-            foo()
-
-        def loss():
-            mul(foo)
-            return jnp.sum(foo.b.value)
-
-        foo = Foo()
-        with bst.StateTraceStack() as trace:
-            grads, loss = bst.augment.grad(loss, foo.states(bst.ParamState), return_value=True)()
-            print(grads)
-            print(loss)
-
-        # print(foo.a.value)
-        # print(foo.b.value)
-        # self.assertTrue(jnp.allclose(foo.a.value, jnp.arange(4) + 1.))
-        # self.assertTrue(jnp.allclose(foo.b.value, jnp.arange(4) * jnp.arange(4)))
-        #
-        # write_state_ids = [id(st) for st in trace.get_write_states()]
-        # read_state_ids = [id(st) for st in trace.get_read_states()]
-        #
-        # assert id(foo.a) in write_state_ids
-        # assert id(foo.b) in write_state_ids
-        #
-        # print(trace.get_write_states())
-        # print(trace.get_read_states())
-
-
+class TestMap(unittest.TestCase):
+    def test_map(self):
+        for dim in [(10,), (10, 10), (10, 10, 10)]:
+            x = bst.random.rand(*dim)
+            r1 = bst.augment.map(lambda a: a + 1, x, batch_size=None)
+            r2 = bst.augment.map(lambda a: a + 1, x, batch_size=2)
+            r3 = bst.augment.map(lambda a: a + 1, x, batch_size=4)
+            r4 = bst.augment.map(lambda a: a + 1, x, batch_size=5)
+            true_r = x + 1
+
+            self.assertTrue(jnp.allclose(r1, true_r))
+            self.assertTrue(jnp.allclose(r2, true_r))
+            self.assertTrue(jnp.allclose(r3, true_r))
+            self.assertTrue(jnp.allclose(r4, true_r))

brainstate/compile/_ad_checkpoint.py

@@ -182,10 +182,12 @@ def checkpoint(
 
     static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
     fun = StatefulFunction(fun, static_argnums=static_argnums)
-    checkpointed_fun = jax.checkpoint(fun.jaxpr_call,
-                                      prevent_cse=prevent_cse,
-                                      policy=policy,
-                                      static_argnums=tuple(i + 1 for i in static_argnums))
+    checkpointed_fun = jax.checkpoint(
+        fun.jaxpr_call,
+        prevent_cse=prevent_cse,
+        policy=policy,
+        static_argnums=tuple(i + 1 for i in static_argnums)
+    )
 
     @functools.wraps(fun.fun)
     def remat_fun(*args, **params):

brainstate/compile/_jit.py

@@ -62,7 +62,6 @@ def _get_jitted_fun(
     **kwargs
 ) -> JittedFunction:
     static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
-    # TODO: add to cache stack for clear_cache
     fun = StatefulFunction(fun, static_argnums=static_argnums, abstracted_axes=abstracted_axes, cache_type='jit')
     jit_fun = jax.jit(fun.jaxpr_call,
                       static_argnums=tuple(i + 1 for i in static_argnums),
@@ -83,10 +82,12 @@ def _get_jitted_fun(
             return fun.fun(*args, **params)
 
         # compile the function and get the state trace
+        # print('Compiling ...')
         state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
         read_state_vals = state_trace.get_read_state_values(True)
 
         # call the jitted function
+        # print('Running ...')
         write_state_vals, outs = jit_fun(state_trace.get_state_values(), *args, **params)
         # write the state values back to the states
         write_back_state_values(state_trace, read_state_vals, write_state_vals)
@@ -127,11 +128,15 @@ def _get_jitted_fun(
         """
         # compile the function and get the state trace
         state_trace = fun.compile_function_and_get_state_trace(*args, **params, return_only_write=True)
-        read_state_vals = state_trace.get_read_state_values(True)
+        read_state_vals = state_trace.get_read_state_values(replace_writen=True)
+        write_state_vals = state_trace.get_write_state_values(replace_read=True)
 
-        # call the jitted function
-        return jit_fun.lower(state_trace.get_state_values(), *args, **params).compile()
+        # compile the model
+        ret = jit_fun.lower(state_trace.get_state_values(), *args, **params).compile()
 
+        # write the state values back to the states
+        write_back_state_values(state_trace, read_state_vals, write_state_vals)
+        return ret
 
     jitted_fun: JittedFunction
 
@@ -290,31 +295,35 @@ def jit(
 
     if isinstance(fun, Missing):
         def wrapper(fun_again: Callable) -> JittedFunction:
-            return _get_jitted_fun(fun_again,
-                                   in_shardings=in_shardings,
-                                   out_shardings=out_shardings,
-                                   static_argnums=static_argnums,
-                                   donate_argnums=donate_argnums,
-                                   donate_argnames=donate_argnames,
-                                   keep_unused=keep_unused,
-                                   device=device,
-                                   backend=backend,
-                                   inline=inline,
-                                   abstracted_axes=abstracted_axes,
-                                   **kwargs)
+            return _get_jitted_fun(
+                fun_again,
+                in_shardings=in_shardings,
+                out_shardings=out_shardings,
+                static_argnums=static_argnums,
+                donate_argnums=donate_argnums,
+                donate_argnames=donate_argnames,
+                keep_unused=keep_unused,
+                device=device,
+                backend=backend,
+                inline=inline,
+                abstracted_axes=abstracted_axes,
+                **kwargs
+            )
 
         return wrapper
 
     else:
-        return _get_jitted_fun(fun,
-                               in_shardings,
-                               out_shardings,
-                               static_argnums,
-                               donate_argnums,
-                               donate_argnames,
-                               keep_unused,
-                               device,
-                               backend,
-                               inline,
-                               abstracted_axes,
-                               **kwargs)
+        return _get_jitted_fun(
+            fun,
+            in_shardings,
+            out_shardings,
+            static_argnums,
+            donate_argnums,
+            donate_argnames,
+            keep_unused,
+            device,
+            backend,
+            inline,
+            abstracted_axes,
+            **kwargs
+        )

brainstate/compile/_loop_collect_return.py

@@ -31,6 +31,7 @@ from ._util import write_back_state_values, wrap_single_fun
 __all__ = [
     # "scan" syntax, which is similar to jax.lax.scan
     'scan', 'checkpointed_scan',
+
     # "for_loop" syntax
     'for_loop', 'checkpointed_for_loop',
 ]
@@ -48,9 +49,9 @@ def _wrap_fun_with_pbar(
     @wraps(fun)
     def new_fun(new_carry, inputs):
         i, old_carry = new_carry
-        old_carry, old_outputs = fun(old_carry, inputs)
-        pbar_runner(unvmap(i, op='none'))
-        return (i + 1, old_carry), old_outputs
+        new_carry, new_outputs = fun(old_carry, inputs)
+        pbar_runner(unvmap(i, op='none'), carry=new_carry, y=new_outputs)
+        return (i + 1, new_carry), new_outputs
 
     return new_fun
 
@@ -476,6 +477,8 @@ def checkpointed_for_loop(
     return ys
 
 
+# This function is adapted from ``while_loop`` in `equinox <https://github.com/patrick-kidger/equinox/blob/main/equinox/internal/_loop/loop.py>`_.
+
 # There's several tricks happening here to work around various limitations of JAX.
 # (Also see https://github.com/google/jax/issues/2139#issuecomment-1039293633)
 # 1. `unvmap_any` prior to using `lax.cond`. JAX has a problem in that vmap-of-cond

brainstate/compile/_loop_no_collection.py

@@ -134,6 +134,8 @@ def bounded_while_loop(
     """
     While loop with a bound on the maximum number of steps.
 
+    This function is adapted from ``while_loop`` in `equinox <https://github.com/patrick-kidger/equinox/blob/main/equinox/internal/_loop/loop.py>`_.
+
     This function is useful when you want to ensure that a while loop terminates
     even if the condition function is never false. The function is implemented
     using a scan operation, so it is reverse-mode differentiable.

brainstate/compile/_make_jaxpr.py

@@ -396,7 +396,7 @@ class StatefulFunction(object):
         # Checking whether the states are returned.
         for leaf in jax.tree.leaves(out):
             if isinstance(leaf, State):
-                leaf._raise_error_with_source_info(ValueError(f"State object is not allowed to be returned: {leaf}"))
+                leaf.raise_error_with_source_info(ValueError(f"State object is not allowed to be returned: {leaf}"))
         return out, state_values
 
     def make_jaxpr(self, *args, return_only_write: bool = False, **kwargs):
@@ -741,11 +741,15 @@ def _make_jaxpr(
         in_type = tuple(jax.util.safe_zip(in_avals, keep_inputs))
         f, out_tree = _flatten_fun(f, in_tree)
         f = annotate(f, in_type)
-        debug_info = pe.debug_info(fun, in_tree, out_tree, True, '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))
-            jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
+            if jax.__version_info__ < (0, 5, 0):
+                jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f, debug_info=debug_info)
+            else:
+                jaxpr, out_type, consts = pe.trace_to_jaxpr_dynamic2(f)
         closed_jaxpr = jax.core.ClosedJaxpr(jaxpr, consts)
         if return_shape:
             out_avals, _ = jax.util.unzip2(out_type)