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

brainstate/augment/_mapping_test.py

@@ -19,6 +19,8 @@ import unittest
 import jax
 import jax.numpy as jnp
 import numpy as np
+from jax import vmap
+from jax.lax import psum, pmean, pmax
 
 import brainstate
 import brainstate.augment
@@ -433,3 +435,163 @@ class TestVMAPNewStatesEdgeCases(unittest.TestCase):
             foo.c = brainstate.State(jnp.arange(3))  # Original expected shape is (4,)
 
         faulty_init()
+
+
+class TestAxisName:
+    def test1(self):
+        def compute_stats_with_axis_name(x):
+            """Compute statistics using named axis operations"""
+            # Sum across the named axis 'batch'
+            total_sum = psum(x, axis_name='batch')
+
+            # Mean across the named axis 'batch'
+            mean_val = pmean(x, axis_name='batch')
+
+            # Max across the named axis 'batch'
+            max_val = pmax(x, axis_name='batch')
+
+            return {
+                'sum': total_sum,
+                'mean': mean_val,
+                'max': max_val,
+                'original': x
+            }
+
+        batch_data = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
+        print("Input batch data:", batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_jax = jax.jit(vmap(compute_stats_with_axis_name, axis_name='batch'))
+        result_jax = vectorized_stats_jax(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats = brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        result = vectorized_stats(batch_data)
+
+        # vmap with axis name 'batch'
+        vectorized_stats_v2 = brainstate.transform.jit(
+            brainstate.transform.vmap(compute_stats_with_axis_name, axis_name='batch')
+        )
+        result_v2 = vectorized_stats_v2(batch_data)
+
+        for key in result_jax.keys():
+            print(f"  {key}: {result_jax[key]}")
+            assert jnp.allclose(result_jax[key], result[key]), f"Mismatch in {key}"
+            assert jnp.allclose(result_jax[key], result_v2[key]), f"Mismatch in {key}"
+
+    def test_nested_vmap(self):
+        def nested_computation(x):
+            """Computation with multiple named axes"""
+            # Sum over 'inner' axis, then mean over 'outer' axis
+            inner_sum = psum(x, axis_name='inner')
+            outer_mean = pmean(inner_sum, axis_name='outer')
+            return outer_mean
+
+        # Create 2D batch data
+        data_2d = jnp.arange(12.0).reshape(3, 4)  # Shape: [outer_batch=3, inner_batch=4]
+        print("Input 2D data shape:", data_2d.shape)
+        print("Input 2D data:\n", data_2d)
+
+        # Nested vmap: first over inner dimension, then outer dimension
+        inner_vmap = vmap(nested_computation, axis_name='inner')
+        nested_vmap = vmap(inner_vmap, axis_name='outer')
+
+        result_2d = nested_vmap(data_2d)
+        print("Result after nested vmap:", result_2d)
+
+        inner_vmap_bst = brainstate.transform.vmap(nested_computation, axis_name='inner')
+        nested_vmap_bst = brainstate.transform.vmap(inner_vmap_bst, axis_name='outer')
+        result_2d_bst = nested_vmap_bst(data_2d)
+        print("Result after nested vmap:", result_2d_bst)
+
+        assert jnp.allclose(result_2d, result_2d_bst)
+
+    def _gradient_averaging_simulation_bst(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def _gradient_averaging_simulation_jax(self):
+        def loss_function(params, x, y):
+            """Simple quadratic loss"""
+            pred = params * x
+            return (pred - y) ** 2
+
+        def compute_gradients_with_averaging(params, batch_x, batch_y):
+            """Compute gradients and average them across the batch"""
+            # Compute per-sample gradients
+            grad_fn = jax.grad(loss_function, argnums=0)
+            per_sample_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+
+            # Average gradients across batch using named axis
+            def average_grads(grads):
+                return pmean(grads, axis_name='batch')
+
+            # Apply averaging with named axis
+            averaged_grads = brainstate.transform.vmap(average_grads, axis_name='batch')(per_sample_grads)
+            return averaged_grads
+
+        # Example data
+        params = 2.0
+        batch_x = jnp.array([1.0, 2.0, 3.0, 4.0])
+        batch_y = jnp.array([2.0, 4.0, 7.0, 8.0])
+
+        print("Parameters:", params)
+        print("Batch X:", batch_x)
+        print("Batch Y:", batch_y)
+
+        # Compute individual gradients first
+        grad_fn = jax.grad(loss_function, argnums=0)
+        individual_grads = brainstate.transform.vmap(grad_fn, in_axes=(None, 0, 0))(params, batch_x, batch_y)
+        print("Individual gradients:", individual_grads)
+
+        # Now compute averaged gradients using axis names
+        averaged_grads = compute_gradients_with_averaging(params, batch_x, batch_y)
+        print("Averaged gradients:", averaged_grads)
+
+        return individual_grads, averaged_grads
+
+    def test_gradient_averaging_simulation(self):
+        individual_grads, averaged_grads = self._gradient_averaging_simulation_bst()
+        individual_grads_jax, averaged_grads_jax = self._gradient_averaging_simulation_jax()
+        assert jnp.allclose(individual_grads, individual_grads_jax)
+        assert jnp.allclose(averaged_grads, averaged_grads_jax)
+
+
+
+

brainstate/compile/_conditions.py

@@ -203,9 +203,9 @@ def ifelse(conditions, branches, *operands, check_cond: bool = True):
     Examples
     --------
 
-    >>> import brainstate as bst
+    >>> import brainstate
     >>> def f(a):
-    >>>    return bst.compile.ifelse(conditions=[a > 10, a > 5, a > 2, a > 0],
+    >>>    return brainstate.compile.ifelse(conditions=[a > 10, a > 5, a > 2, a > 0],
     >>>                               branches=[lambda: 1,
     >>>                                         lambda: 2,
     >>>                                         lambda: 3,

brainstate/compile/_make_jaxpr.py

@@ -352,6 +352,13 @@ class StatefulFunction(PrettyObject):
             cache_key = default_cache_key
         return self.get_state_trace(cache_key).get_write_states()
 
+    def _check_input(self, x):
+        if isinstance(x, State):
+            raise ValueError(
+                'Inputs for brainstate transformations cannot be an instance of State. '
+                f'But we got {x}'
+            )
+
     def get_arg_cache_key(self, *args, **kwargs) -> Tuple:
         """
         Get the static arguments from the arguments.
@@ -370,22 +377,35 @@ class StatefulFunction(PrettyObject):
                     static_args.append(arg)
                 else:
                     dyn_args.append(arg)
-            dyn_args = jax.tree.map(shaped_abstractify, jax.tree.leaves(dyn_args))
+            dyn_args = jax.tree.map(shaped_abstractify, dyn_args)
             static_kwargs, dyn_kwargs = [], []
             for k, v in kwargs.items():
                 if k in self.static_argnames:
                     static_kwargs.append((k, v))
                 else:
                     dyn_kwargs.append((k, jax.tree.map(shaped_abstractify, v)))
-            return tuple([tuple(static_args), tuple(dyn_args), tuple(static_kwargs), tuple(dyn_kwargs)])
+
+            static_args = make_hashable(tuple(static_args))
+            dyn_args = make_hashable(tuple(dyn_args))
+            static_kwargs = make_hashable(static_kwargs)
+            dyn_kwargs = make_hashable(dyn_kwargs)
+
+            cache_key = (static_args, dyn_args, static_kwargs, dyn_kwargs)
         elif self.cache_type is None:
             num_arg = len(args)
             static_args = tuple(args[i] for i in self.static_argnums if i < num_arg)
             static_kwargs = tuple((k, v) for k, v in kwargs.items() if k in self.static_argnames)
-            return tuple([static_args, static_kwargs])
+
+            # Make everything hashable
+            static_args = make_hashable(static_args)
+            static_kwargs = make_hashable(static_kwargs)
+
+            cache_key = (static_args, static_kwargs)
         else:
             raise ValueError(f"Invalid cache type: {self.cache_type}")
 
+        return cache_key
+
     def compile_function_and_get_states(self, *args, **kwargs) -> Tuple[State, ...]:
         """
         Compile the function, and get the states that are read and written by this function.
@@ -480,6 +500,9 @@ class StatefulFunction(PrettyObject):
         # static args
         cache_key = self.get_arg_cache_key(*args, **kwargs)
 
+        # check input types
+        jax.tree.map(self._check_input, (args, kwargs), is_leaf=lambda x: isinstance(x, State))
+
         if cache_key not in self._cached_state_trace:
             try:
                 # jaxpr
@@ -637,15 +660,15 @@ def make_jaxpr(
     instead give a few examples.
 
     >>> import jax
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>>
     >>> def f(x): return jax.numpy.sin(jax.numpy.cos(x))
     >>> print(f(3.0))
     -0.83602
-    >>> jaxpr, states = bst.compile.make_jaxpr(f)(3.0)
+    >>> jaxpr, states = brainstate.compile.make_jaxpr(f)(3.0)
     >>> jaxpr
     { lambda ; a:f32[]. let b:f32[] = cos a; c:f32[] = sin b in (c,) }
-    >>> jaxpr, states = bst.compile.make_jaxpr(jax.grad(f))(3.0)
+    >>> jaxpr, states = brainstate.compile.make_jaxpr(jax.grad(f))(3.0)
     >>> jaxpr
     { lambda ; a:f32[]. let
         b:f32[] = cos a
@@ -844,3 +867,33 @@ def _make_jaxpr(
     if hasattr(fun, "__name__"):
         make_jaxpr_f.__name__ = f"make_jaxpr({fun.__name__})"
     return make_jaxpr_f
+
+
+def make_hashable(obj):
+    """Convert a pytree into a hashable representation."""
+    if isinstance(obj, (list, tuple)):
+        return tuple(make_hashable(item) for item in obj)
+    elif isinstance(obj, dict):
+        return tuple(sorted((k, make_hashable(v)) for k, v in obj.items()))
+    elif isinstance(obj, set):
+        return frozenset(make_hashable(item) for item in obj)
+    elif hasattr(obj, '__dict__'):  # Handle custom objects
+        return (
+            obj.__class__.__name__,
+            tuple(
+                sorted(
+                    (k, make_hashable(v))
+                    for k, v in obj.__dict__.items()
+                    if not k.startswith('_')
+                )
+            )
+        )
+    else:
+        # # Use JAX's tree_util for any other pytree structures
+        # try:
+        #     leaves, treedef = jax.tree_util.tree_flatten(obj)
+        #     hashable_leaves = tuple(make_hashable(leaf) for leaf in leaves)
+        #     return (str(treedef), hashable_leaves)
+        # except:
+        #     # Assume obj is already hashable
+        return obj

brainstate/compile/_progress_bar.py

@@ -53,7 +53,7 @@ class ProgressBar(object):
 
     .. code-block:: python
 
-                a = bst.State(1.)
+                a = brainstate.State(1.)
                 def loop_fn(x):
                     a.value = x.value + 1.
                     return jnp.sum(x ** 2)
@@ -61,7 +61,7 @@ class ProgressBar(object):
                 pbar = ProgressBar(desc=("Running {i} iterations, loss = {loss}",
                                          lambda i_carray_y: {"i": i_carray_y["i"], "loss": i_carray_y["y"]}))
 
-                bst.compile.for_loop(loop_fn, xs, pbar=pbar)
+                brainstate.compile.for_loop(loop_fn, xs, pbar=pbar)
 
     In this example, ``"i"`` denotes the iteration number and ``"loss"`` is computed from the output,
     the ``"carry"`` is the dynamic state in the loop, for example ``a.value`` in this case.

brainstate/environ.py

@@ -76,9 +76,9 @@ def context(**kwargs):
 
     For instance::
 
-    >>> import brainstate as bst
-    >>> with bst.environ.context(dt=0.1) as env:
-    ...     dt = bst.environ.get('dt')
+    >>> import brainstate as brainstate
+    >>> with brainstate.environ.context(dt=0.1) as env:
+    ...     dt = brainstate.environ.get('dt')
     ...     print(env)
 
     """
@@ -424,10 +424,10 @@ def dftype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default floating data type is ``np.float32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.dftype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.dftype())
     >>> print(a.dtype)
 
     Returns
@@ -448,10 +448,10 @@ def ditype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default integer data type is ``np.int32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.ditype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.ditype())
     >>> print(a.dtype)
     int32
 
@@ -474,10 +474,10 @@ def dutype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default unsigned integer data type is ``np.uint32``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...   a = np.zeros(1, dtype=bst.environ.dutype())
+    >>> with brainstate.environ.context(precision=32):
+    ...   a = np.zeros(1, dtype=brainstate.environ.dutype())
     >>> print(a.dtype)
     uint32
 
@@ -499,10 +499,10 @@ def dctype() -> DTypeLike:
 
     For example, if the precision is set to 32, the default complex data type is ``np.complex64``.
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> import numpy as np
-    >>> with bst.environ.context(precision=32):
-    ...    a = np.zeros(1, dtype=bst.environ.dctype())
+    >>> with brainstate.environ.context(precision=32):
+    ...    a = np.zeros(1, dtype=brainstate.environ.dctype())
     >>> print(a.dtype)
     complex64
 
@@ -529,19 +529,19 @@ def register_default_behavior(key: str, behavior: Callable, replace_if_exist: bo
 
     For example, you can register a default behavior for the key 'dt' by::
 
-    >>> import brainstate as bst
+    >>> import brainstate as brainstate
     >>> def dt_behavior(dt):
     ...     print(f'Set the default dt to {dt}.')
     ...
-    >>> bst.environ.register_default_behavior('dt', dt_behavior)
+    >>> brainstate.environ.register_default_behavior('dt', dt_behavior)
 
     Then, when you set the default dt by `brainstate.environ.set(dt=0.1)`, the behavior
     `dt_behavior` will be called with
     `dt_behavior(0.1)`.
 
-    >>> bst.environ.set(dt=0.1)
+    >>> brainstate.environ.set(dt=0.1)
     Set the default dt to 0.1.
-    >>> with bst.environ.context(dt=0.2):
+    >>> with brainstate.environ.context(dt=0.2):
     ...     pass
     Set the default dt to 0.2.
     Set the default dt to 0.1.

brainstate/functional/_activations_test.py

@@ -70,39 +70,39 @@ class NNFunctionsTest(jtu.JaxTestCase):
             check_dtypes=False)
 
     #   def testSquareplusGrad(self):
-    #     check_grads(bst.functional.squareplus, (1e-8,), order=4,
+    #     check_grads(brainstate.functional.squareplus, (1e-8,), order=4,
     #                 )
 
     #   def testSquareplusGradZero(self):
-    #     check_grads(bst.functional.squareplus, (0.,), order=1,
+    #     check_grads(brainstate.functional.squareplus, (0.,), order=1,
     #                 )
 
     #   def testSquareplusGradNegInf(self):
-    #     check_grads(bst.functional.squareplus, (-float('inf'),), order=1,
+    #     check_grads(brainstate.functional.squareplus, (-float('inf'),), order=1,
     #                 )
 
     #   def testSquareplusGradNan(self):
-    #     check_grads(bst.functional.squareplus, (float('nan'),), order=1,
+    #     check_grads(brainstate.functional.squareplus, (float('nan'),), order=1,
     #                 )
 
     #   @parameterized.parameters([float] + jtu.dtypes.floating)
     #   def testSquareplusZero(self, dtype):
-    #     self.assertEqual(dtype(1), bst.functional.squareplus(dtype(0), dtype(4)))
+    #     self.assertEqual(dtype(1), brainstate.functional.squareplus(dtype(0), dtype(4)))
     #
     # def testMishGrad(self):
-    #   check_grads(bst.functional.mish, (1e-8,), order=4,
+    #   check_grads(brainstate.functional.mish, (1e-8,), order=4,
     #               )
     #
     # def testMishGradZero(self):
-    #   check_grads(bst.functional.mish, (0.,), order=1,
+    #   check_grads(brainstate.functional.mish, (0.,), order=1,
     #               )
     #
     # def testMishGradNegInf(self):
-    #   check_grads(bst.functional.mish, (-float('inf'),), order=1,
+    #   check_grads(brainstate.functional.mish, (-float('inf'),), order=1,
     #               )
     #
     # def testMishGradNan(self):
-    #   check_grads(bst.functional.mish, (float('nan'),), order=1,
+    #   check_grads(brainstate.functional.mish, (float('nan'),), order=1,
     #               )
 
     @parameterized.parameters([float] + jtu.dtypes.floating)
@@ -137,7 +137,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
         self.assertAllClose(brainstate.functional.sparse_sigmoid(0.), .5, check_dtypes=False)
 
     #   def testSquareplusValue(self):
-    #     val = bst.functional.squareplus(1e3)
+    #     val = brainstate.functional.squareplus(1e3)
     #     self.assertAllClose(val, 1e3, check_dtypes=False, atol=1e-3)
 
     def testMishValue(self):
@@ -177,7 +177,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
          brainstate.functional.softplus,
          brainstate.functional.sparse_plus,
          brainstate.functional.sigmoid,
-         #  bst.functional.squareplus,
+         #  brainstate.functional.squareplus,
          brainstate.functional.mish)))
     def testDtypeMatchesInput(self, dtype, fn):
         x = jnp.zeros((), dtype=dtype)
@@ -306,7 +306,7 @@ class NNFunctionsTest(jtu.JaxTestCase):
 
     def testCustomJVPLeak2(self):
         # https://github.com/google/jax/issues/8171
-        # The above test uses jax.bst.functional.sigmoid, as in the original #8171, but that
+        # The above test uses jax.brainstate.functional.sigmoid, as in the original #8171, but that
         # function no longer actually has a custom_jvp! So we inline the old def.
 
         @jax.custom_jvp