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

brainstate/nn/_collective_ops.py

@@ -16,14 +16,18 @@
 from __future__ import annotations
 
 from collections import namedtuple
-from typing import Dict, Callable, TypeVar
 
 import jax
+from typing import (
+    Callable, TypeVar, Tuple, Any, Dict
+)
 
 from brainstate._state import catch_new_states
 from brainstate._utils import set_module_as
+from brainstate.augment import vmap, vmap_new_states
 from brainstate.graph import nodes
-from brainstate.util._filter import Filter
+from brainstate.random import set_key, split_key
+from brainstate.typing import Filter
 from ._module import Module
 
 # the maximum order
@@ -35,8 +39,16 @@ StateLoadResult = namedtuple('StateLoadResult', ['missing_keys', 'unexpected_key
 T = TypeVar('T', bound=Module)
 
 __all__ = [
-    'MAX_ORDER', 'call_order', 'init_all_states', 'reset_all_states',
-    'load_all_states', 'save_all_states', 'assign_state_values',
+    'MAX_ORDER',
+    'call_order',
+    'call_all_functions',
+    'vmap_call_all_functions',
+    'init_all_states',
+    'vmap_init_all_states',
+    'reset_all_states',
+    'load_all_states',
+    'save_all_states',
+    'assign_state_values',
 ]
 
 
@@ -76,77 +88,368 @@ def call_order(level: int = 0, check_order_boundary: bool = True):
 
 
 @set_module_as('brainstate.nn')
-def init_all_states(
+def call_all_functions(
     target: T,
-    *args,
-    exclude: Filter = None,
-    **kwargs
+    fun_name: str,
+    args: Tuple[Any, ...] | Any = (),
+    kwargs: Dict[str, Any] | None = None,
+    node_to_exclude: Filter = None,
+    fun_if_not_exist: str = 'raise',
 ) -> T:
     """
-    Collectively initialize states of all children nodes in the given target.
-
-    Args:
-      target: The target Module.
-      exclude: The filter to exclude some nodes.
-      tag: The tag for the new states.
-      args: The positional arguments for the initialization, which will be passed to the `init_state` method
-        of each node.
-      kwargs: The keyword arguments for the initialization, which will be passed to the `init_state` method
-        of each node.
-
+    Call a specified function on all nodes of a target module, respecting call order if defined.
+
+    This function iterates through all nodes of the target module, calling a specified function
+    on each node. It respects the call order of functions if defined, and provides options for
+    handling cases where the specified function does not exist on a node.
+
+    Parameters:
+    -----------
+    target : T
+        The target module on which to call functions.
+    fun_name : str
+        The name of the function to call on each node.
+    args : Tuple[Any, ...] | Any, optional
+        Positional arguments to pass to the called function. Default is an empty tuple.
+    kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to pass to the called function. Default is None.
+    node_to_exclude : Filter, optional
+        A filter function to exclude certain nodes from the function call.
+    fun_if_not_exist : str, optional
+        Specifies behavior when the function doesn't exist on a node. Options are:
+
+        - 'raise': Raise an exception (default)
+        - 'pass' or 'none': Skip the node and continue
+        
     Returns:
-      The target Module.
+    --------
+    T
+        The target module after calling the specified function on all applicable nodes.
+
+    Raises:
+    -------
+    AssertionError
+        If fun_name is not a string or kwargs is not a dictionary.
+    ValueError
+        If fun_if_not_exist is not one of the allowed values.
+    AttributeError
+        If the specified function doesn't exist on a node and fun_if_not_exist is 'raise'.
     """
+    assert isinstance(fun_name, str), f'fun_name must be a string, but got {fun_name}.'
 
-    # node that has `call_order` decorated
-    nodes_with_order = []
+    args = (args,) if not isinstance(args, tuple) else args
+    kwargs = kwargs or {}
+    assert isinstance(kwargs, dict), f'kwargs must be a dict, but got {kwargs}.'
 
-    nodes_ = nodes(target).filter(Module)
-    if exclude is not None:
-        nodes_ = nodes_ - nodes_.filter(exclude)
+    all_nodes = nodes(target).filter(Module)
+    if node_to_exclude is not None:
+        all_nodes -= all_nodes.filter(node_to_exclude)
 
-    # reset node whose `init_state` has no `call_order`
-    for node in list(nodes_.values()):
-        if hasattr(node.init_state, 'call_order'):
+    nodes_with_order = []
+    for node in all_nodes.values():
+        try:
+            fun = getattr(node, fun_name)
+        except AttributeError as e:
+            if fun_if_not_exist == 'raise':
+                raise
+            elif fun_if_not_exist in ('pass', 'none'):
+                continue
+            else:
+                raise ValueError(
+                    f'fun_if_not_exist must be one of ["raise", "pass", "none"], but got {fun_if_not_exist}.')
+
+        assert callable(fun), f'{fun_name} must be a callable function, but got {fun}.'
+        if hasattr(fun, 'call_order'):
             nodes_with_order.append(node)
         else:
-            node.init_state(*args, **kwargs)
+            fun(*args, **kwargs)
+
+    for node in sorted(nodes_with_order, key=lambda x: getattr(x, fun_name).call_order):
+        getattr(node, fun_name)(*args, **kwargs)
+
+    return target
+
+
+def vmap_call_all_functions(
+    target: T,
+    fun_name: str,
+    args: Tuple[Any, ...] | Any = (),
+    kwargs: Dict[str, Any] | None = None,
+    axis_size: int = None,
+    node_to_exclude: Filter = None,
+    tag: str | None = None,
+    fun_if_not_exist: str = 'raise',
+) -> T:
+    """
+    Apply vectorized mapping (vmap) to call a specified function on all nodes of a target module.
+
+    This function vectorizes the process of calling a specified function across multiple instances
+    of the target module, effectively batching the operation.
+
+    Parameters:
+    -----------
+    target : T
+        The target module on which to call functions.
+    fun_name : str
+        The name of the function to call on each node.
+    args : Tuple[Any, ...] | Any, optional
+        Positional arguments to pass to the called function. Default is an empty tuple.
+    kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to pass to the called function. Default is None.
+    axis_size : int, optional
+        The size of the batch axis for vmap. Must be a positive integer.
+    node_to_exclude : Filter, optional
+        A filter function to exclude certain nodes from the function call.
+    tag : str | None, optional
+        A tag to be used for catching new states.
+    fun_if_not_exist : str, optional
+        Specifies behavior when the function doesn't exist on a node. Options are:
+
+        - 'raise': Raise an exception (default)
+        - 'pass' or 'none': Skip the node and continue
+
+    Returns:
+    --------
+    T
+        The target module after applying the vectorized function call on all applicable nodes.
 
-    # reset the node's states with `call_order`
-    for node in sorted(nodes_with_order, key=lambda x: x.init_state.call_order):
-        node.init_state(*args, **kwargs)
+    Raises:
+    -------
+    AssertionError
+        If axis_size is not specified or is not a positive integer.
+    """
+    assert axis_size is not None and axis_size > 0, f"axis_size must be a positive integer, got {axis_size}"
+
+    if not isinstance(args, tuple):
+        args = (args,)
+    kwargs = kwargs or {}
+    assert isinstance(kwargs, dict), f'kwargs must be a dict, but got {kwargs}.'
+
+    @vmap(out_axes=0, axis_size=axis_size)
+    def vmapped_fn(key):
+        set_key(key)
+        with catch_new_states(tag) as inner_catcher:
+            call_all_functions(
+                target,
+                fun_name=fun_name,
+                args=args,
+                kwargs=kwargs,
+                node_to_exclude=node_to_exclude,
+                fun_if_not_exist=fun_if_not_exist
+            )
+        values = inner_catcher.get_state_values()
+        return values
+
+    with catch_new_states(tag) as outer_catcher:
+        values = vmapped_fn(split_key(axis_size))
+        states = outer_catcher.get_states()
+    for state, value in zip(states, values):
+        state.value = value
 
     return target
 
 
 @set_module_as('brainstate.nn')
-def reset_all_states(target: Module, *args, **kwargs) -> Module:
+def init_all_states(
+    target: T,
+    init_args: Tuple[Any, ...] | Any = (),
+    init_kwargs: Dict[str, Any] | None = None,
+    node_to_exclude: Filter = None,
+) -> T:
     """
-    Collectively reset states of all children nodes in the given target.
+    Initialize all states for the given target module and its submodules.
 
-    Args:
-      target: The target Module.
+    This function initializes the states of the target module and all its submodules,
+    respecting any call order decorators that may be present on the init_state methods.
 
-    Returns:
-      The target Module.
+    Parameters
+    ----------
+    target : T
+        The target module whose states are to be initialized.
+    init_args : Tuple[Any, ...] | Any, optional
+        Positional arguments to be passed to each init_state method.
+        If a single non-tuple argument is provided, it will be wrapped in a tuple.
+    init_kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to be passed to each init_state method.
+        If None, an empty dictionary will be used.
+    node_to_exclude : Filter, optional
+        A filter function or predicate to exclude certain nodes from initialization.
+
+    Returns
+    -------
+    T
+        The target module with all states initialized.
+
+    Raises
+    ------
+    AssertionError
+        If init_kwargs is provided but is not a dictionary.
     """
+    return call_all_functions(target, 'init_state', init_args, init_kwargs, node_to_exclude)
 
-    nodes_with_order = []
 
-    # reset node whose `init_state` has no `call_order`
-    for path, node in nodes(target).filter(Module).items():
-        if hasattr(node.reset_state, 'call_order'):
-            nodes_with_order.append(node)
-        else:
-            node.reset_state(*args, **kwargs)
+@set_module_as('brainstate.nn')
+def vmap_init_all_states(
+    target: T,
+    init_args: Tuple[Any, ...] | Any = (),
+    init_kwargs: Dict[str, Any] | None = None,
+    axis_size: int = None,
+    node_to_exclude: Filter = None,
+    state_to_exclude: Filter = None,
+    state_tag: str | None = None,
+) -> T:
+    """
+    Initialize all vmap states for the given target module.
+
+    This function applies vectorized mapping (vmap) to initialize states across multiple
+    instances of the target module, effectively batching the initialization process.
+
+    Parameters:
+    -----------
+    target : T
+        The target module whose states are to be initialized.
+    init_args : Tuple[Any, ...] | Any, optional
+        Positional arguments to be passed to the init_all_states function. Default is an empty tuple.
+    init_kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to be passed to the init_all_states function. Default is None.
+    axis_size : int, optional
+        The size of the batch axis for vmap. This must be specified and should be greater than 0.
+    node_to_exclude : Filter, optional
+        A filter to exclude certain nodes from initialization.
+    state_tag : str | None, optional
+        A tag to be used for catching new states.
 
-    # reset the node's states
-    for node in sorted(nodes_with_order, key=lambda x: x.reset_state.call_order):
-        node.reset_state(*args, **kwargs)
+    Returns:
+    --------
+    T
+        The target module with initialized states.
 
+    Raises:
+    -------
+    AssertionError
+        If axis_size is not specified or is not greater than 0.
+        If init_kwargs is not a dictionary.
+    """
+
+    # return vmap_call_all_functions(
+    #     target,
+    #     'init_state',
+    #     args=init_args,
+    #     kwargs=init_kwargs,
+    #     axis_size=axis_size,
+    #     node_to_exclude=node_to_exclude,
+    #     tag=tag,
+    # )
+
+    def init_fn():
+        init_all_states(
+            target,
+            init_args=init_args,
+            init_kwargs=init_kwargs,
+            node_to_exclude=node_to_exclude,
+        )
+        return
+
+    vmap_new_states(init_fn, state_tag=state_tag, axis_size=axis_size, state_to_exclude=state_to_exclude)()
     return target
 
 
+@set_module_as('brainstate.nn')
+def reset_all_states(
+    target: T,
+    reset_args: Tuple[Any, ...] | Any = (),
+    reset_kwargs: Dict[str, Any] | None = None,
+    node_to_exclude: Filter = None,
+) -> T:
+    """
+    Reset all states for the given target module and its submodules.
+
+    This function resets the states of the target module and all its submodules,
+    respecting any call order decorators that may be present on the reset_state methods.
+
+    Parameters
+    ----------
+    target : T
+        The target module whose states are to be reset.
+    reset_args : Tuple[Any, ...] | Any, optional
+        Positional arguments to be passed to each reset_state method.
+        If a single non-tuple argument is provided, it will be wrapped in a tuple.
+    reset_kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to be passed to each reset_state method.
+        If None, an empty dictionary will be used.
+    node_to_exclude : Filter, optional
+        A filter function or predicate to exclude certain nodes from reset.
+
+    Returns
+    -------
+    T
+        The target module with all states reset.
+
+    Raises
+    ------
+    AssertionError
+        If init_kwargs is provided but is not a dictionary.
+    """
+    return call_all_functions(
+        target,
+        fun_name='reset_state',
+        args=reset_args,
+        kwargs=reset_kwargs,
+        node_to_exclude=node_to_exclude
+    )
+
+
+def vmap_reset_all_states(
+    target: T,
+    reset_args: Tuple[Any, ...] | Any = (),
+    reset_kwargs: Dict[str, Any] | None = None,
+    axis_size: int = None,
+    node_to_exclude: Filter = None,
+    tag: str | None = None,
+) -> T:
+    """
+    Reset all vmap states for the given target module.
+
+    This function applies vectorized mapping (vmap) to reset states across multiple
+    instances of the target module, effectively batching the reset process.
+
+    Parameters:
+    -----------
+    target : T
+        The target module whose states are to be reset.
+    reset_args : Tuple[Any, ...] | Any, optional
+        Positional arguments to be passed to the reset_all_states function. Default is an empty tuple.
+    reset_kwargs : Dict[str, Any] | None, optional
+        Keyword arguments to be passed to the reset_all_states function. Default is None.
+    axis_size : int, optional
+        The size of the batch axis for vmap. This must be specified and should be greater than 0.
+    node_to_exclude : Filter, optional
+        A filter to exclude certain nodes from reset.
+    tag : str | None, optional
+        A tag to be used for catching new states.
+
+    Returns:
+    --------
+    T
+        The target module with reset states.
+
+    Raises:
+    -------
+    AssertionError
+        If axis_size is not specified or is not greater than 0.
+        If reset_kwargs is not a dictionary.
+    """
+    return vmap_call_all_functions(
+        target,
+        fun_name='reset_state',
+        args=reset_args,
+        kwargs=reset_kwargs,
+        axis_size=axis_size,
+        node_to_exclude=node_to_exclude,
+        tag=tag,
+    )
+
+
 @set_module_as('brainstate.nn')
 def load_all_states(target: Module, state_dict: Dict, **kwargs):
     """

brainstate/nn/_collective_ops_test.py

@@ -0,0 +1,36 @@
+# Copyright 2025 BDP 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+
+import brainstate as bst
+
+
+class Test_vmap_init_all_states:
+
+    def test_vmap_init_all_states(self):
+        gru = bst.nn.GRUCell(1, 2)
+        bst.nn.vmap_init_all_states(gru, axis_size=10)
+        print(gru)
+
+    def test_vmap_init_all_states_v2(self):
+        @bst.compile.jit
+        def init():
+            gru = bst.nn.GRUCell(1, 2)
+            bst.nn.vmap_init_all_states(gru, axis_size=10)
+            print(gru)
+
+        init()

brainstate/nn/_common.py

@@ -0,0 +1,194 @@
+# Copyright 2025 BDP 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.
+# ==============================================================================
+
+# -*- coding: utf-8 -*-
+
+from __future__ import annotations
+
+from collections import defaultdict
+
+from typing import Any, Sequence, Hashable, Dict
+
+from brainstate import environ
+from brainstate.augment._mapping import vmap
+from brainstate.typing import Filter
+from ._module import Module
+
+AxisName = Hashable
+
+__all__ = [
+    'EnvironContext',
+    'Vmap',
+]
+
+
+class EnvironContext(Module):
+    """
+    A wrapper class that provides an environment context for a given layer.
+
+    This class allows execution of a layer within a specific environment context,
+    which can be useful for controlling the execution environment of neural network layers.
+
+    This class is equivalent to the following code snippet:
+
+    ```python
+
+    import brainstate
+
+    with brainstate.environ.context(**context):
+        result = layer(*args, **kwargs)
+
+    ```
+
+    Attributes:
+        layer (Module): The layer to be executed within the environment context.
+        context (dict): The environment context parameters.
+    """
+
+    def __init__(self, layer: Module, **context):
+        """
+        Initialize the EnvironContext.
+
+        Args:
+            layer (Module): The layer to be wrapped with the environment context.
+            **context: Arbitrary keyword arguments representing the environment context parameters.
+        """
+        super().__init__()
+
+        assert isinstance(layer, Module), 'The layer must be an instance of Module.'
+        self.layer = layer
+        self.context = context
+
+    def update(self, *args, **kwargs):
+        """
+        Execute the wrapped layer within the specified environment context.
+
+        Args:
+            *args: Variable length argument list to be passed to the wrapped layer.
+            **kwargs: Arbitrary keyword arguments to be passed to the wrapped layer.
+
+        Returns:
+            The result of executing the wrapped layer within the environment context.
+        """
+        with environ.context(**self.context):
+            return self.layer(*args, **kwargs)
+
+    def add_context(self, **context):
+        """
+        Add additional environment context parameters to the existing context.
+
+        Args:
+            **context: Arbitrary keyword arguments representing the additional environment context parameters.
+        """
+        self.context.update(context)
+
+
+def _filter_states(
+    module: Module,
+    filters: Filter | Dict[Filter, int],
+) -> Dict:
+    if filters is None:
+        filtered_states = None
+    elif isinstance(filters, dict):
+        in_states_filter = defaultdict(list)
+        for filter_, axis in filters:
+            assert isinstance(axis, int), 'The value of in_states must be the map axis, which should be an integer.'
+            in_states_filter[axis].append(filter_)
+        filtered_states = module.states(*in_states_filter.values())
+        in_states_axis = tuple(in_states_filter.keys())
+        filtered_states = {axis: states for axis, states in zip(in_states_axis, filtered_states)}
+    else:
+        filtered_states = module.states(filters)
+    return filtered_states
+
+
+class Vmap(Module):
+    """
+    A class that applies vectorized mapping (vmap) to a given module.
+
+    This class wraps a module and applies vectorized mapping to its execution,
+    allowing for efficient parallel processing across specified axes.
+
+    Attributes:
+        module (Module): The module to be vmapped.
+        in_axes (int | None | Sequence[Any]): Specifies how to map over inputs.
+        out_axes (Any): Specifies how to map over outputs.
+        vmap_states (Filter | Dict[Filter, int]): Specifies which states to vmap and on which axes.
+        vmap_out_states (Filter | Dict[Filter, int]): Specifies which output states to vmap and on which axes.
+        axis_name (AxisName | None): Name of the axis being mapped over.
+        axis_size (int | None): Size of the axis being mapped over.
+    """
+
+    def __init__(
+        self,
+        module: Module,
+        in_axes: int | None | Sequence[Any] = 0,
+        out_axes: Any = 0,
+        vmap_states: Filter | Dict[Filter, int] = None,
+        vmap_out_states: Filter | Dict[Filter, int] = None,
+        axis_name: AxisName | None = None,
+        axis_size: int | None = None,
+    ):
+        """
+        Initialize the Vmap instance.
+
+        Args:
+            module (Module): The module to be vmapped.
+            in_axes (int | None | Sequence[Any], optional): Specifies how to map over inputs. Defaults to 0.
+            out_axes (Any, optional): Specifies how to map over outputs. Defaults to 0.
+            vmap_states (Filter | Dict[Filter, int], optional): Specifies which states to vmap and on which axes. Defaults to None.
+            vmap_out_states (Filter | Dict[Filter, int], optional): Specifies which output states to vmap and on which axes. Defaults to None.
+            axis_name (AxisName | None, optional): Name of the axis being mapped over. Defaults to None.
+            axis_size (int | None, optional): Size of the axis being mapped over. Defaults to None.
+        """
+        super().__init__()
+
+        # parameters
+        self.in_axes = in_axes
+        self.out_axes = out_axes
+        self.axis_name = axis_name
+        self.axis_size = axis_size
+        assert isinstance(module, Module), 'The module must be an instance of Module.'
+        self.module = module
+        vmap_states = _filter_states(module, vmap_states)
+        vmap_out_states = _filter_states(module, vmap_out_states)
+
+        @vmap(
+            in_axes=in_axes,
+            out_axes=out_axes,
+            in_states=vmap_states,
+            out_states=vmap_out_states,
+            axis_name=axis_name,
+            axis_size=axis_size,
+        )
+        def vmap_run(*args, **kwargs):
+            vmap_states_ = vmap_states
+            return module(*args, **kwargs)
+
+        # vmapped module
+        self.vmapped_fn = vmap_run
+
+    def update(self, *args, **kwargs):
+        """
+        Execute the vmapped module with the given arguments.
+
+        Args:
+            *args: Variable length argument list to be passed to the vmapped module.
+            **kwargs: Arbitrary keyword arguments to be passed to the vmapped module.
+
+        Returns:
+            The result of executing the vmapped module.
+        """
+        return self.vmapped_fn(*args, **kwargs)

brainstate/nn/_dyn_impl/_dynamics_neuron.py

@@ -17,10 +17,9 @@
 
 from __future__ import annotations
 
-from typing import Callable, Optional
-
 import brainunit as u
 import jax
+from typing import Callable, Optional
 
 from brainstate import init, surrogate, environ
 from brainstate._state import HiddenState, ShortTermState