brainstate 0.1.10__py2.py3-none-any.whl → 0.2.0__py2.py3-none-any.whl

brainstate/transform/_mapping_test.py

@@ -0,0 +1,194 @@
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+
+import unittest
+
+import jax
+import jax.numpy as jnp
+from jax import vmap
+from jax.lax import psum, pmean, pmax
+
+import brainstate
+import brainstate.transform
+from brainstate._error import BatchAxisError
+
+
+
+class TestMap(unittest.TestCase):
+    def test_map(self):
+        for dim in [(10,), (10, 10), (10, 10, 10)]:
+            x = brainstate.random.rand(*dim)
+            r1 = brainstate.transform.map(lambda a: a + 1, x, batch_size=None)
+            r2 = brainstate.transform.map(lambda a: a + 1, x, batch_size=2)
+            r3 = brainstate.transform.map(lambda a: a + 1, x, batch_size=4)
+            r4 = brainstate.transform.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))
+
+
+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
+

brainstate/{compile → transform}/_progress_bar.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -41,42 +41,95 @@ class ProgressBar(object):
 
     The message displayed in the progress bar can be customized by the following two methods:
 
-    1. By passing a string to the `desc` argument. For example:
+    1. By passing a string to the `desc` argument.
+    2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
+       function should take a dictionary as input and return a dictionary. The returned dictionary
+       will be used to format the string.
+
+    In the second case, ``"i"`` denotes the iteration number and other keys can be computed from the
+    loop outputs and carry values.
+
+    Parameters
+    ----------
+    freq : int, optional
+        The frequency at which to print the progress bar. If not specified, the progress
+        bar will be printed every 5% of the total iterations.
+    count : int, optional
+        The number of times to print the progress bar. If not specified, the progress
+        bar will be printed every 5% of the total iterations. Cannot be used together with `freq`.
+    desc : str or tuple, optional
+        A description of the progress bar. If not specified, a default message will be
+        displayed. Can be either a string or a tuple of (format_string, format_function).
+    **kwargs
+        Additional keyword arguments to pass to the progress bar.
+
+    Examples
+    --------
+    Basic usage with default description:
 
     .. code-block:: python
 
-            ProgressBar(desc="Running 1000 iterations")
+        >>> import brainstate
+        >>> import jax.numpy as jnp
+        >>>
+        >>> def loop_fn(x):
+        ...     return x ** 2
+        >>>
+        >>> xs = jnp.arange(100)
+        >>> pbar = brainstate.transform.ProgressBar()
+        >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
 
-    2. By passing a tuple with a string and a callable function to the `desc` argument. The callable
-       function should take a dictionary as input and return a dictionary. The returned dictionary
-       will be used to format the string. For example:
+    With custom description string:
+
+    .. code-block:: python
+
+        >>> pbar = brainstate.transform.ProgressBar(desc="Running 1000 iterations")
+        >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
+
+    With frequency control:
 
     .. code-block:: python
 
-                a = brainstate.State(1.)
-                def loop_fn(x):
-                    a.value = x.value + 1.
-                    return jnp.sum(x ** 2)
+        >>> # Update every 10 iterations
+        >>> pbar = brainstate.transform.ProgressBar(freq=10)
+        >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
+        >>>
+        >>> # Update exactly 20 times during execution
+        >>> pbar = brainstate.transform.ProgressBar(count=20)
+        >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
 
-                pbar = ProgressBar(desc=("Running {i} iterations, loss = {loss}",
-                                         lambda i_carray_y: {"i": i_carray_y["i"], "loss": i_carray_y["y"]}))
+    With dynamic description based on loop variables:
 
-                brainstate.compile.for_loop(loop_fn, xs, pbar=pbar)
+    .. code-block:: python
 
-    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.
+        >>> state = brainstate.State(1.0)
+        >>>
+        >>> def loop_fn(x):
+        ...     state.value += x
+        ...     loss = jnp.sum(x ** 2)
+        ...     return loss
+        >>>
+        >>> def format_desc(data):
+        ...     return {"i": data["i"], "loss": data["y"], "state": data["carry"]}
+        >>>
+        >>> pbar = brainstate.transform.ProgressBar(
+        ...     desc=("Iteration {i}, loss = {loss:.4f}, state = {state:.2f}", format_desc)
+        ... )
+        >>> results = brainstate.transform.for_loop(loop_fn, xs, pbar=pbar)
+
+    With scan function:
 
+    .. code-block:: python
 
-    Args:
-        freq: The frequency at which to print the progress bar. If not specified, the progress
-            bar will be printed every 5% of the total iterations.
-        count: The number of times to print the progress bar. If not specified, the progress
-            bar will be printed every 5% of the total iterations.
-        desc: A description of the progress bar. If not specified, a default message will be
-            displayed.
-        kwargs: Additional keyword arguments to pass to the progress bar.
+        >>> def scan_fn(carry, x):
+        ...     new_carry = carry + x
+        ...     return new_carry, new_carry ** 2
+        >>>
+        >>> init_carry = 0.0
+        >>> pbar = brainstate.transform.ProgressBar(freq=5)
+        >>> final_carry, ys = brainstate.transform.scan(scan_fn, init_carry, xs, pbar=pbar)
     """
-    __module__ = "brainstate.compile"
+    __module__ = "brainstate.transform"
 
     def __init__(
         self,
@@ -141,7 +194,7 @@ class ProgressBar(object):
 
 
 class ProgressBarRunner(object):
-    __module__ = "brainstate.compile"
+    __module__ = "brainstate.transform"
 
     def __init__(
         self,

brainstate/{augment → transform}/_random.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -16,6 +16,7 @@
 import functools
 from typing import Callable, Sequence, Union
 
+from brainstate._utils import set_module_as
 from brainstate.random import DEFAULT, RandomState
 from brainstate.typing import Missing
 from brainstate.util import PrettyObject
@@ -27,47 +28,66 @@ __all__ = [
 
 class RngRestore(PrettyObject):
     """
-    Backup and restore the random state of a sequence of RandomState instances.
+    Manage backing up and restoring multiple random states.
 
-    This class provides functionality to save the current state of multiple
-    RandomState instances and later restore them to their saved states.
+    Parameters
+    ----------
+    rngs : Sequence[RandomState]
+        Sequence of :class:`~brainstate.random.RandomState` instances whose
+        states should be captured and restored.
+
+    Attributes
+    ----------
+    rngs : Sequence[RandomState]
+        Managed random-state instances.
+    rng_keys : list
+        Cached keys captured by :meth:`backup` until :meth:`restore` runs.
 
-    Attributes:
-        rngs (Sequence[RandomState]): A sequence of RandomState instances to manage.
-        rng_keys (list): A list to store the backed up random keys.
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> rng = brainstate.random.RandomState(0)
+        >>> restorer = brainstate.transform.RngRestore([rng])
+        >>> restorer.backup()
+        >>> _ = rng.random()
+        >>> restorer.restore()
     """
+    __module__ = 'brainstate.transform'
 
     def __init__(self, rngs: Sequence[RandomState]):
         """
-        Initialize the RngRestore instance.
+        Initialize a restorer for the provided random states.
 
-        Args:
-            rngs (Sequence[RandomState]): A sequence of RandomState instances
-                whose states will be managed.
+        Parameters
+        ----------
+        rngs : Sequence[RandomState]
+            Random states that will be backed up and restored.
         """
         self.rngs: Sequence[RandomState] = rngs
         self.rng_keys = []
 
     def backup(self):
         """
-        Backup the current random key of the RandomState instances.
+        Cache the current key for each managed random state.
 
-        This method saves the current value (state) of each RandomState
-        instance in the rngs sequence.
+        Notes
+        -----
+        The cached keys persist until :meth:`restore` is called, after which the
+        internal cache is cleared.
         """
         self.rng_keys = [rng.value for rng in self.rngs]
 
     def restore(self):
         """
-        Restore the random key of the RandomState instances.
+        Restore each random state to the cached key.
 
-        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.
+        Raises
+        ------
+        ValueError
+            Raised when the number of stored keys does not match ``rngs``.
         """
         if len(self.rng_keys) != len(self.rngs):
             raise ValueError('The number of random keys does not match the number of random states.')
@@ -95,50 +115,50 @@ def _rng_backup(
     return wrapper
 
 
+@set_module_as('brainstate.transform')
 def restore_rngs(
     fn: Callable = Missing(),
     rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
 ) -> Callable:
     """
-    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.
+    Decorate a function so specified random states are restored after execution.
 
     Parameters
     ----------
     fn : Callable, optional
-        The function to be wrapped. If not provided, the decorator can be used
-        with parameters.
+        Function to wrap. When omitted, :func:`restore_rngs` returns a decorator
+        preconfigured with ``rngs``.
     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.
+        Random states whose keys should be backed up before running ``fn`` and
+        restored afterwards. Defaults to :data:`brainstate.random.DEFAULT`.
 
     Returns
     -------
     Callable
-        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`.
+        Wrapped callable that restores the random state or a partially applied
+        decorator depending on how :func:`restore_rngs` is used.
 
     Raises
     ------
     AssertionError
-        If `rngs` is not a RandomState instance or a sequence of RandomState instances.
+        If ``rngs`` is neither a :class:`~brainstate.random.RandomState` instance nor
+        a sequence of such 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()
+    .. code-block:: python
+
+        >>> import brainstate
+        >>>
+        >>> rng = brainstate.random.RandomState(0)
+        >>>
+        >>> @brainstate.transform.restore_rngs(rngs=rng)
+        ... def sample_pair():
+        ...     first = rng.random()
+        ...     second = rng.random()
+        ...     return first, second
+        >>>
+        >>> assert sample_pair()[0] == sample_pair()[0]
     """
     if isinstance(fn, Missing):
         return functools.partial(restore_rngs, rngs=rngs)

brainstate/{compile → transform}/_unvmap.py

@@ -1,4 +1,4 @@
-# Copyright 2024 BDP Ecosystem Limited. All Rights Reserved.
+# Copyright 2024 BrainX Ecosystem Limited. All Rights Reserved.
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -27,8 +27,39 @@ __all__ = [
 ]
 
 
-@set_module_as('brainstate.augment')
+@set_module_as('brainstate.transform')
 def unvmap(x, op: str = 'any'):
+    """
+    Remove a leading vmap dimension by aggregating batched values.
+
+    Parameters
+    ----------
+    x : Any
+        Value produced inside a :func:`jax.vmap`-transformed function.
+    op : {'all', 'any', 'none', 'max'}, default='any'
+        Reduction to apply across the vmapped axis. ``'none'`` returns ``x`` without
+        reduction, while ``'max'`` computes the maximum element.
+
+    Returns
+    -------
+    Any
+        Result of applying the requested reduction with vmap metadata removed.
+
+    Raises
+    ------
+    ValueError
+        If ``op`` is not one of ``'all'``, ``'any'``, ``'none'``, or ``'max'``.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate
+        >>>
+        >>> xs = jnp.array([[True, False], [True, True]])
+        >>> brainstate.transform.unvmap(xs, op='all')
+    """
     if op == 'all':
         return unvmap_all(x)
     elif op == 'any':
@@ -47,7 +78,29 @@ unvmap_all_p = Primitive("unvmap_all")
 
 
 def unvmap_all(x):
-    """As `jnp.all`, but ignores batch dimensions."""
+    """
+    Evaluate :func:`jax.numpy.all` while ignoring vmapped batch dimensions.
+
+    Parameters
+    ----------
+    x : Any
+        Input array or pytree produced under :func:`jax.vmap`.
+
+    Returns
+    -------
+    jax.Array
+        Scalar boolean result of ``jnp.all(x)``.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate
+        >>>
+        >>> values = jnp.array([[True, False], [True, True]])
+        >>> brainstate.transform.unvmap(values, op='all')
+    """
     return unvmap_all_p.bind(x)
 
 
@@ -78,7 +131,29 @@ unvmap_any_p = Primitive("unvmap_any")
 
 
 def unvmap_any(x):
-    """As `jnp.any`, but ignores batch dimensions."""
+    """
+    Evaluate :func:`jax.numpy.any` while ignoring vmapped batch dimensions.
+
+    Parameters
+    ----------
+    x : Any
+        Input array or pytree produced under :func:`jax.vmap`.
+
+    Returns
+    -------
+    jax.Array
+        Scalar boolean result of ``jnp.any(x)``.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate
+        >>>
+        >>> values = jnp.array([[False, False], [False, True]])
+        >>> brainstate.transform.unvmap(values, op='any')
+    """
     return unvmap_any_p.bind(x)
 
 
@@ -109,7 +184,29 @@ unvmap_max_p = Primitive("unvmap_max")
 
 
 def unvmap_max(x):
-    """As `jnp.max`, but ignores batch dimensions."""
+    """
+    Evaluate :func:`jax.numpy.max` while ignoring vmapped batch dimensions.
+
+    Parameters
+    ----------
+    x : Any
+        Input array or pytree produced under :func:`jax.vmap`.
+
+    Returns
+    -------
+    jax.Array
+        Scalar containing the maximum value of ``x`` with the same dtype.
+
+    Examples
+    --------
+    .. code-block:: python
+
+        >>> import jax.numpy as jnp
+        >>> import brainstate
+        >>>
+        >>> values = jnp.array([[1.0, 2.0], [0.5, 3.5]])
+        >>> brainstate.transform.unvmap(values, op='max')
+    """
     return unvmap_max_p.bind(x)