brainstate 0.1.0.post20250212__py2.py3-none-any.whl → 0.1.0.post20250217__py2.py3-none-any.whl

brainstate/augment/_mapping_test.py

@@ -15,12 +15,14 @@
 
 from __future__ import annotations
 
-import unittest
-
+import jax
 import jax.numpy as jnp
+import numpy as np
+import unittest
 
 import brainstate as bst
 from brainstate.augment._mapping import BatchAxisError
+from brainstate.augment._mapping import _remove_axis
 
 
 class TestVmap(unittest.TestCase):
@@ -99,6 +101,27 @@ class TestVmap(unittest.TestCase):
         )
         print(bst.random.DEFAULT)
 
+    def test_vmap_with_random_v3(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):
+                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
+        )()
+        print(bst.random.DEFAULT)
+
     def test_vmap_with_random_2(self):
         class Model(bst.nn.Module):
             def __init__(self):
@@ -114,22 +137,11 @@ class TestVmap(unittest.TestCase):
                 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,
+            out_states=model.c
         )(
             bst.random.split_key(5)
         )
@@ -154,24 +166,17 @@ class TestVmap(unittest.TestCase):
         print(model.weight.value_call(jnp.shape))
         print(model.weight.value)
 
-    def test_vmap_model(self):
-        model = bst.nn.Linear(2, 3)
-        model_id = id(model)
-        weight_id = id(model.weight)
-        print(id(model), id(model.weight))
-        x = jnp.ones((5, 2))
+    def test_vmap_states_and_input_1(self):
+        gru = bst.nn.GRUCell(2, 3)
+        gru.init_state(5)
 
-        @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)
-            print(id(model), id(model.weight))
-            return model(x)
+        @bst.augment.vmap(in_states=gru.states(bst.HiddenState))
+        def forward(x):
+            return gru(x)
 
-        y = forward(model, x)
-        print(y.shape)
-        print(model.weight.value_call(jnp.shape))
-        print(model.weight.value)
+        xs = bst.random.randn(5, 2)
+        y = forward(xs)
+        self.assertTrue(y.shape == (5, 3))
 
     def test_vmap_jit(self):
         class Foo(bst.nn.Module):
@@ -249,6 +254,16 @@ class TestVmap(unittest.TestCase):
         print(trace.get_write_states())
         print(trace.get_read_states())
 
+    def test_auto_rand_key_split(self):
+        def f():
+            return bst.random.rand(1)
+
+        res = bst.augment.vmap(f, axis_size=10)()
+        self.assertTrue(jnp.all(~(res[0] == res[1:])))
+
+        res2 = jax.vmap(f, axis_size=10)()
+        self.assertTrue(jnp.all((res2[0] == res2[1:])))
+
 
 class TestMap(unittest.TestCase):
     def test_map(self):
@@ -264,3 +279,72 @@ class TestMap(unittest.TestCase):
             self.assertTrue(jnp.allclose(r2, true_r))
             self.assertTrue(jnp.allclose(r3, true_r))
             self.assertTrue(jnp.allclose(r4, true_r))
+
+
+class TestRemoveAxis:
+
+    def test_remove_axis_2d_array_axis_0(self):
+        input_array = np.array([[1, 2, 3], [4, 5, 6]])
+        expected_output = np.array([1, 2, 3])
+
+        result = _remove_axis(input_array, axis=0)
+
+        np.testing.assert_array_equal(result, expected_output)
+
+    def test_remove_axis_3d_array(self):
+        # Create a 3D array
+        x = np.arange(24).reshape((2, 3, 4))
+
+        # Remove axis 1
+        result = _remove_axis(x, axis=1)
+
+        # Expected result: a 2D array with shape (2, 4)
+        expected = x[:, 0, :]
+
+        np.testing.assert_array_equal(result, expected)
+        assert result.shape == (2, 4)
+
+    def test_remove_axis_1d_array(self):
+        # Create a 1D array
+        x = np.array([1, 2, 3, 4, 5])
+
+        # Remove axis 0 (the only axis in a 1D array)
+        result = _remove_axis(x, axis=0)
+
+        # Check that the result is a scalar (0D array) and equal to the first element
+        assert np.isscalar(result), "Result should be a scalar"
+        assert result == 1, "Result should be equal to the first element of the input array"
+
+    def test_remove_axis_out_of_bounds(self):
+        x = jnp.array([[1, 2], [3, 4]])
+        with unittest.TestCase().assertRaises(IndexError):
+            _remove_axis(x, axis=2)
+
+    def test_remove_axis_negative(self):
+        x = jnp.array([[[1, 2], [3, 4]], [[5, 6], [7, 8]]])
+        result = _remove_axis(x, -1)
+        expected = jnp.array([[1, 3], [5, 7]])
+        np.testing.assert_array_equal(result, expected)
+
+    def test_remove_axis_with_nan_and_inf(self):
+        x = jnp.array([[1.0, jnp.nan, 3.0], [4.0, 5.0, jnp.inf]])
+        result = _remove_axis(x, axis=0)
+        expected = jnp.array([1.0, jnp.nan, 3.0])
+        np.testing.assert_array_equal(result, expected)
+        assert jnp.isnan(result[1])
+
+    def test_remove_axis_different_dtypes(self):
+        # Test with integer array
+        int_array = jnp.array([[1, 2, 3], [4, 5, 6]])
+        int_result = _remove_axis(int_array, 0)
+        assert jnp.array_equal(int_result, jnp.array([1, 2, 3]))
+
+        # Test with float array
+        float_array = jnp.array([[1.1, 2.2, 3.3], [4.4, 5.5, 6.6]])
+        float_result = _remove_axis(float_array, 1)
+        assert jnp.allclose(float_result, jnp.array([1.1, 4.4]))
+
+        # Test with complex array
+        complex_array = jnp.array([[1 + 1j, 2 + 2j], [3 + 3j, 4 + 4j]])
+        complex_result = _remove_axis(complex_array, 0)
+        assert jnp.allclose(complex_result, jnp.array([1 + 1j, 2 + 2j]))

brainstate/augment/_random.py

@@ -20,34 +20,62 @@ from typing import Callable, Sequence, Union
 
 from brainstate.random import DEFAULT, RandomState
 from brainstate.typing import Missing
+from brainstate.util import PrettyObject
 
 __all__ = [
     'restore_rngs'
 ]
 
 
-class RngRestore:
+class RngRestore(PrettyObject):
     """
     Backup and restore the random state of a sequence of RandomState instances.
+
+    This class provides functionality to save the current state of multiple
+    RandomState instances and later restore them to their saved states.
+
+    Attributes:
+        rngs (Sequence[RandomState]): A sequence of RandomState instances to manage.
+        rng_keys (list): A list to store the backed up random keys.
     """
 
     def __init__(self, rngs: Sequence[RandomState]):
+        """
+        Initialize the RngRestore instance.
+
+        Args:
+            rngs (Sequence[RandomState]): A sequence of RandomState instances
+                whose states will be managed.
+        """
         self.rngs: Sequence[RandomState] = rngs
         self.rng_keys = []
 
     def backup(self):
         """
         Backup the current random key of the RandomState instances.
+
+        This method saves the current value (state) of each RandomState
+        instance in the rngs sequence.
         """
         self.rng_keys = [rng.value for rng in self.rngs]
 
     def restore(self):
         """
         Restore the random key of the RandomState instances.
+
+        This method restores each RandomState instance to its previously
+        saved state. It raises an error if the number of saved keys doesn't
+        match the number of RandomState instances.
+
+        Raises:
+            ValueError: If the number of saved random keys does not match
+                the number of RandomState instances.
         """
+        if len(self.rng_keys) != len(self.rngs):
+            raise ValueError('The number of random keys does not match the number of random states.')
         for rng, key in zip(self.rngs, self.rng_keys):
             rng.restore_value(key)
-        self.rng_keys = []
+        self.rng_keys.clear()
 
 
 def _rng_backup(
@@ -74,19 +102,45 @@ def restore_rngs(
     rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
 ) -> Callable:
     """
-    Backup the current random state and restore it after the function call.
+    Decorator to backup and restore the random state before and after a function call.
+
+    This function can be used as a decorator or called directly. It ensures that the
+    random state of the specified RandomState instances is preserved across function calls,
+    which is useful for maintaining reproducibility in stochastic operations.
 
     Parameters
     ----------
     fn : Callable, optional
-        The function to be wrapped.
-    rngs : Union[RandomState, Sequence[RandomState]]
-        The random state to be backed up and restored. If not provided, the default RandomState instance will be used.
+        The function to be wrapped. If not provided, the decorator can be used
+        with parameters.
+    rngs : Union[RandomState, Sequence[RandomState]], optional
+        The random state(s) to be backed up and restored. This can be a single
+        RandomState instance or a sequence of RandomState instances. If not provided,
+        the default RandomState instance will be used.
 
     Returns
     -------
     Callable
-        The wrapped function.
+        If `fn` is provided, returns the wrapped function that will backup the
+        random state before execution and restore it afterwards.
+        If `fn` is not provided, returns a partial function that can be used as
+        a decorator with the specified `rngs`.
+
+    Raises
+    ------
+    AssertionError
+        If `rngs` is not a RandomState instance or a sequence of RandomState instances.
+
+    Examples
+    --------
+    >>> @restore_rngs
+    ... def my_random_function():
+    ...     return random.random()
+
+    >>> rng = RandomState(42)
+    >>> @restore_rngs(rngs=rng)
+    ... def another_random_function():
+    ...     return rng.random()
     """
     if isinstance(fn, Missing):
         return functools.partial(restore_rngs, rngs=rngs)

brainstate/compile/_ad_checkpoint.py

@@ -16,9 +16,8 @@
 from __future__ import annotations
 
 import functools
-from typing import Callable, Tuple, Union
-
 import jax
+from typing import Callable, Tuple, Union
 
 from brainstate.typing import Missing
 from ._make_jaxpr import StatefulFunction, _ensure_index_tuple
@@ -181,7 +180,7 @@ def checkpoint(
         return lambda f: checkpoint(f, prevent_cse=prevent_cse, policy=policy, static_argnums=static_argnums)
 
     static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
-    fun = StatefulFunction(fun, static_argnums=static_argnums)
+    fun = StatefulFunction(fun, static_argnums=static_argnums, name='checkpoint')
     checkpointed_fun = jax.checkpoint(
         fun.jaxpr_call,
         prevent_cse=prevent_cse,

brainstate/compile/_conditions.py

@@ -15,11 +15,10 @@
 
 from __future__ import annotations
 
-from collections.abc import Callable, Sequence
-
 import jax
 import jax.numpy as jnp
 import numpy as np
+from collections.abc import Callable, Sequence
 
 from brainstate._utils import set_module_as
 from ._error_if import jit_error_if
@@ -94,8 +93,8 @@ def cond(pred, true_fun: Callable, false_fun: Callable, *operands):
             return false_fun(*operands)
 
     # evaluate jaxpr
-    stateful_true = StatefulFunction(true_fun).make_jaxpr(*operands)
-    stateful_false = StatefulFunction(false_fun).make_jaxpr(*operands)
+    stateful_true = StatefulFunction(true_fun, name='cond:true').make_jaxpr(*operands)
+    stateful_false = StatefulFunction(false_fun, name='conda:false').make_jaxpr(*operands)
 
     # state trace and state values
     state_trace = stateful_true.get_state_trace() + stateful_false.get_state_trace()
@@ -174,7 +173,7 @@ def switch(index, branches: Sequence[Callable], *operands):
         return branches[int(index)](*operands)
 
     # evaluate jaxpr
-    wrapped_branches = [StatefulFunction(branch) for branch in branches]
+    wrapped_branches = [StatefulFunction(branch, name='switch') for branch in branches]
     for wrapped_branch in wrapped_branches:
         wrapped_branch.make_jaxpr(*operands)
 

brainstate/compile/_conditions_test.py

@@ -14,10 +14,9 @@
 # ==============================================================================
 from __future__ import annotations
 
-import unittest
-
 import jax
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 

brainstate/compile/_error_if.py

@@ -16,11 +16,10 @@
 from __future__ import annotations
 
 import functools
+import jax
 from functools import partial
 from typing import Callable, Union
 
-import jax
-
 from brainstate._utils import set_module_as
 from ._unvmap import unvmap
 

brainstate/compile/_error_if_test.py

@@ -15,11 +15,10 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax
 import jax.numpy as jnp
 import jaxlib.xla_extension
+import unittest
 
 import brainstate as bst
 

brainstate/compile/_jit.py

@@ -16,12 +16,11 @@
 from __future__ import annotations
 
 import functools
-from collections.abc import Iterable, Sequence
-from typing import (Any, Callable, Union)
-
 import jax
+from collections.abc import Iterable, Sequence
 from jax._src import sharding_impls
 from jax.lib import xla_client as xc
+from typing import (Any, Callable, Union)
 
 from brainstate._utils import set_module_as
 from brainstate.typing import Missing
@@ -62,19 +61,27 @@ def _get_jitted_fun(
     **kwargs
 ) -> JittedFunction:
     static_argnums = _ensure_index_tuple(tuple() if static_argnums is None else static_argnums)
-    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),
-                      donate_argnums=donate_argnums,
-                      donate_argnames=donate_argnames,
-                      keep_unused=keep_unused,
-                      device=device,
-                      backend=backend,
-                      inline=inline,
-                      in_shardings=in_shardings,
-                      out_shardings=out_shardings,
-                      abstracted_axes=abstracted_axes,
-                      **kwargs)
+    fun = StatefulFunction(
+        fun,
+        static_argnums=static_argnums,
+        abstracted_axes=abstracted_axes,
+        cache_type='jit',
+        name='jit'
+    )
+    jit_fun = jax.jit(
+        fun.jaxpr_call,
+        static_argnums=tuple(i + 1 for i in static_argnums),
+        donate_argnums=donate_argnums,
+        donate_argnames=donate_argnames,
+        keep_unused=keep_unused,
+        device=device,
+        backend=backend,
+        inline=inline,
+        in_shardings=in_shardings,
+        out_shardings=out_shardings,
+        abstracted_axes=abstracted_axes,
+        **kwargs
+    )
 
     @functools.wraps(fun.fun)
     def jitted_fun(*args, **params):

brainstate/compile/_jit_test.py

@@ -15,9 +15,8 @@
 
 from __future__ import annotations
 
-import unittest
-
 import jax.numpy as jnp
+import unittest
 
 import brainstate as bst
 

brainstate/compile/_loop_collect_return.py

@@ -16,11 +16,11 @@
 from __future__ import annotations
 
 import math
-from functools import wraps
-from typing import Callable, Optional, TypeVar, Tuple, Any
 
 import jax
 import jax.numpy as jnp
+from functools import wraps
+from typing import Callable, Optional, TypeVar, Tuple, Any
 
 from brainstate._utils import set_module_as
 from ._make_jaxpr import StatefulFunction
@@ -209,7 +209,7 @@ def scan(
     # ------------------------------ #
     xs_avals = [jax.core.get_aval(x) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
-    stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
+    stateful_fun = StatefulFunction(f, name='scan').make_jaxpr(init, xs_tree.unflatten(x_avals))
     state_trace = stateful_fun.get_state_trace()
     all_writen_state_vals = state_trace.get_write_state_values(True)
     all_read_state_vals = state_trace.get_read_state_values(True)
@@ -217,12 +217,20 @@ def scan(
 
     # scan
     init = (all_writen_state_vals, init)
-    (all_writen_state_vals, carry), ys = jax.lax.scan(wrapped_f,
-                                                      init,
-                                                      xs,
-                                                      length=length,
-                                                      reverse=reverse,
-                                                      unroll=unroll)
+    (
+        (
+            all_writen_state_vals,
+            carry
+        ),
+        ys
+    ) = jax.lax.scan(
+        wrapped_f,
+        init,
+        xs,
+        length=length,
+        reverse=reverse,
+        unroll=unroll
+    )
     # assign the written state values and restore the read state values
     write_back_state_values(state_trace, all_read_state_vals, all_writen_state_vals)
     # carry
@@ -305,7 +313,7 @@ def checkpointed_scan(
     # evaluate jaxpr
     xs_avals = [jax.core.get_aval(x) for x in xs_flat]
     x_avals = [jax.core.mapped_aval(length, 0, aval) for aval in xs_avals]
-    stateful_fun = StatefulFunction(f).make_jaxpr(init, xs_tree.unflatten(x_avals))
+    stateful_fun = StatefulFunction(f, name='checkpoint_scan').make_jaxpr(init, xs_tree.unflatten(x_avals))
     state_trace = stateful_fun.get_state_trace()
     # get all states
     been_written = state_trace.been_writen

brainstate/compile/_loop_collect_return_test.py

@@ -14,10 +14,10 @@
 # ==============================================================================
 
 from __future__ import annotations
-import unittest
 
 import jax.numpy as jnp
 import numpy as np
+import unittest
 
 import brainstate as bst
 

brainstate/compile/_loop_no_collection.py

@@ -16,9 +16,9 @@
 from __future__ import annotations
 
 import math
-from typing import Any, Callable, TypeVar
 
 import jax
+from typing import Any, Callable, TypeVar
 
 from brainstate._utils import set_module_as
 from ._loop_collect_return import _bounded_while_loop
@@ -103,8 +103,8 @@ def while_loop(
             pass
 
     # evaluate jaxpr
-    stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
-    stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+    stateful_cond = StatefulFunction(cond_fun, name='while:cond').make_jaxpr(init_val)
+    stateful_body = StatefulFunction(body_fun, name='while:body').make_jaxpr(init_val)
     if len(stateful_cond.get_write_states()) != 0:
         raise ValueError("while_loop: cond_fun should not have any write states.")
 
@@ -162,8 +162,8 @@ def bounded_while_loop(
         return init_val
 
     # evaluate jaxpr
-    stateful_cond = StatefulFunction(cond_fun).make_jaxpr(init_val)
-    stateful_body = StatefulFunction(body_fun).make_jaxpr(init_val)
+    stateful_cond = StatefulFunction(cond_fun, name='bounded_while:cond').make_jaxpr(init_val)
+    stateful_body = StatefulFunction(body_fun, name='bounded_while:body').make_jaxpr(init_val)
     if len(stateful_cond.get_write_states()) != 0:
         raise ValueError("while_loop: cond_fun should not have any write states.")