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

brainstate/augment/_mapping.py

@@ -15,152 +15,152 @@
 
 from __future__ import annotations
 
-import dataclasses
 import functools
-from typing import Any, TypeVar, Callable, Hashable, Sequence, Iterable, Mapping, Tuple, Union, Optional
+from typing import Any, TypeVar, Callable, Hashable, Sequence, Iterable, Tuple, Union, Optional, Dict, List
 
 import jax
+from jax.interpreters.batching import BatchTracer
 
-from brainstate.graph import (NodeStates, graph_to_tree, tree_to_graph, update_context)
-from brainstate.graph._graph_convert import clear_non_graph_nodes
+from brainstate._state import State, StateTraceStack
+from brainstate.compile._loop_collect_return import scan
 from brainstate.random import DEFAULT, RandomState
-from brainstate.typing import Missing, Filter
-from brainstate.util import NestedDict
+from brainstate.typing import Missing
+from brainstate.util import NestedDict, BrainStateError
 from ._random import restore_rngs
 
 __all__ = [
-    'StateAxes',
     'vmap',
     'pmap',
+    'map',
 ]
 
 AxisName = Hashable
 F = TypeVar("F", bound=Callable)
-Index = int
-Carry = TypeVar("Carry")
+AxisToState = Dict[int, List[State]]
+StateToAxis = Dict[State, int]
 
 
-class StateAxes:
-    """
-    A class to represent the axes of a state.
-
-    This class is used to control how graph nodes like Modules are vectorized or
-    parallelized by specifying the axes to be applied to substates of the graph
-    node given a Filter.
-
-    Args:
-        filter_axes: A mapping from filters to axes. The axes can be an index, a carry or None.
-
-    """
-
-    def __init__(
-        self,
-        filter_axes: Union[Mapping[Filter, Index | Carry | None], Iterable[Tuple[Filter, Index | Carry | None]]],
-    ):
-        iterable = filter_axes.items() if isinstance(filter_axes, Mapping) else filter_axes
-        self._filters = tuple(filter_ for filter_, _ in iterable)
-        self._axes = tuple(axis for _, axis in iterable)
-
-    @property
-    def filters(self) -> Tuple[Filter, ...]:
-        return self._filters
-
-    @property
-    def axes(self) -> Tuple[Index | Carry | None, ...]:
-        return self._axes
-
-    def __repr__(self):
-        return f'StateAxes({dict(self.items())})'
-
-    def items(self):
-        return zip(self.filters, self.axes)
-
-    def __eq__(self, other):
-        return isinstance(other, StateAxes) and self.filters == other.filters and self.axes == other.axes
-
-    def __hash__(self):
-        return hash((self.filters, self.axes))
+class BatchAxisError(BrainStateError):
+    pass
 
 
-def _map_split_fn(ctx, path, prefix, x):
-    if isinstance(prefix, StateAxes):
-        return NodeStates.from_split(*ctx.treefy_split(x, *prefix.filters), metadata=prefix)
-    return NodeStates.from_split(*ctx.treefy_split(x), metadata=prefix)
-
-
-@dataclasses.dataclass(eq=False)
-class MapFn:
-    f: Callable[..., Any]
-    in_axes: Any
-    out_axes: Any
-    ctxtag: str
-
-    def __post_init__(self):
-        functools.update_wrapper(self, self.f)
-
-    def __call__(self, *pure_args: Tuple[Any, ...]):
-        # pytree to graph
-        args = tree_to_graph(pure_args, ctxtag=self.ctxtag)
-
-        # call the function
-        out = self.f(*args)
-
-        # graph to pytree
-        args_out = clear_non_graph_nodes(args)
-        pure_args_out, pure_out = graph_to_tree(
-            (args_out, out),
-            prefix=(self.in_axes, self.out_axes),
-            split_fn=_map_split_fn,
-            ctxtag=self.ctxtag,
+def _flatten_in_out_states(
+    in_states: Dict[int, Dict] | Any = None,
+) -> Tuple[AxisToState, StateToAxis]:
+    if in_states is None:
+        return dict(), dict()
+    if isinstance(in_states, dict):
+        keys = tuple(in_states.keys())
+        values = tuple(in_states.values())
+        is_axis_in_states = (
+            all([isinstance(key, int) for key in keys]) and
+            all([isinstance(value, dict) for value in values])
         )
-        return pure_args_out, pure_out
-
-
-def _map_transform(
-    ctxtag,
-    transform,
+    else:
+        is_axis_in_states = False
+    if is_axis_in_states:
+        axis_to_states = {key: list(value.values()) for key, value in in_states.items()}
+        state_to_axis = {}
+        for key, value in in_states.items():
+            for state in value.values():
+                state_to_axis[state] = key
+        return axis_to_states, state_to_axis
+    else:
+        in_states = jax.tree.leaves(in_states)
+        axis_to_states = {0: list(in_states)}
+        state_to_axis = {state: 0 for state in in_states}
+        return axis_to_states, state_to_axis
+
+
+def _vmap_transform(
     f: F,
     *,
-    in_axes: Optional[int | Sequence[Any]] = 0,
+    in_axes: int | None | Sequence[Any] = 0,
     out_axes: Any = 0,
+    in_states: Dict[int, Dict] | Any | None = None,
+    out_states: Dict[int, Dict] | Any | None = None,
     rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
     **transform_kwargs,
 ):
-    # jax in axes
-    jax_in_axes = jax.tree.map(
-        lambda x: NodeStates.from_prefixes(x.axes, metadata=x) if isinstance(x, StateAxes) else x,
-        in_axes,
-    )
-
-    # jax out axes
-    jax_out_axes = jax.tree.map(
-        lambda x: NodeStates.from_prefixes(x.axes, metadata=x) if isinstance(x, StateAxes) else x,
-        out_axes,
-    )
-
-    # mapped function
-    mapped_fn = transform(
-        MapFn(f, in_axes, out_axes, ctxtag),
-        in_axes=jax_in_axes,
-        out_axes=(jax_in_axes, jax_out_axes),
-        **transform_kwargs
-    )
+    axis_to_in_states, in_state_to_axis = _flatten_in_out_states(in_states)
+    axis_to_out_states, out_state_to_axis = _flatten_in_out_states(out_states)
+    for _in_state, _axis in in_state_to_axis.items():
+        if _in_state in out_state_to_axis:
+            _out_axis = out_state_to_axis[_in_state]
+            if _out_axis != _axis:
+                _in_state.raise_error_with_source_info(
+                    BatchAxisError(
+                        f"State {_in_state} has been mapped to axis {_axis} in in_states, "
+                        f"However, it is mapped to axis {_out_axis} in out_states."
+                    )
+                )
+        else:
+            out_state_to_axis[_in_state] = _axis
+            if _axis not in axis_to_out_states:
+                axis_to_out_states[_axis] = []
+            axis_to_out_states[_axis].append(_in_state)
+    if isinstance(rngs, RandomState):
+        rngs = (rngs,)
+    rng_ids = set([id(rng) for rng in rngs])
 
     @functools.wraps(f)
-    @restore_rngs(rngs=rngs)  # restore the random key of default random number generator
-    @update_context(ctxtag)
-    def map_wrapper(*args):
-        # graph to pytree
-        pure_args = graph_to_tree(args, prefix=in_axes, split_fn=_map_split_fn, ctxtag=ctxtag)
-
-        # vmap with pytree
-        pure_args_out, pure_out = mapped_fn(*pure_args)
-
-        # pytree to graph
-        _args_out, out = tree_to_graph((pure_args_out, pure_out), ctxtag=ctxtag)
-        return out
+    def new_fn(in_states_, args):
+        # restore state values
+        for i, states in enumerate(axis_to_in_states.values()):
+            for state, state_val in zip(states, in_states_[i]):
+                state.restore_value(state_val)
 
-    return map_wrapper  # type: ignore
+        # call the function
+        with StateTraceStack() as stack:
+            outs = f(*args)
+
+        # analyze
+        for state in stack.get_write_states():
+            leaves = jax.tree.leaves(state.value)
+            if isinstance(leaves[0], BatchTracer) and state not in out_state_to_axis:
+                if isinstance(state, RandomState) and id(state) in rng_ids:
+                    continue
+                state.raise_error_with_source_info(
+                    BatchAxisError(
+                        f"The value of State {state} is batched, but it is not in the out_states."
+                    )
+                )
+
+        out_states_ = [
+            [state.value for state in states]
+            for axis, states in axis_to_out_states.items()
+        ]
+        return out_states_, outs
+
+    def vmapped_fn(*args):
+        # vmapping
+        in_state_vals = [
+            [st.value for st in states]
+            for axis, states in axis_to_in_states.items()
+        ]
+        in_axes_st = list(axis_to_in_states.keys())
+        out_axes_st = list(axis_to_out_states.keys())
+        if len(in_axes_st) == 0:
+            in_axes_st = 0
+        if len(out_axes_st) == 0:
+            out_axes_st = 0
+        out_state_vals, outs = restore_rngs(
+            jax.vmap(
+                new_fn,
+                in_axes=(in_axes_st, in_axes),
+                out_axes=(out_axes_st, out_axes),
+                **transform_kwargs
+            ),
+            rngs=rngs
+        )(in_state_vals, args)
+
+        # restore mapped state values
+        for i, states in enumerate(axis_to_out_states.values()):
+            for state, st_val in zip(states, out_state_vals[i]):
+                state.restore_value(st_val)
+        return outs
+
+    return vmapped_fn
 
 
 def vmap(
@@ -172,6 +172,8 @@ def vmap(
     axis_size: int | None = None,
     spmd_axis_name: AxisName | tuple[AxisName, ...] | None = None,
     # brainstate specific arguments
+    in_states: Dict[int, Dict] | Any | None = None,
+    out_states: Dict[int, Dict] | Any | None = None,
     rngs: Union[RandomState, Sequence[RandomState]] = DEFAULT,
 ) -> F | Callable[[F], F]:
     """
@@ -183,103 +185,38 @@ def vmap(
 
     More information please see `jax.vmap <https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html>`__.
 
-
     These are several example usage::
 
         >>> import brainstate as bst
         >>> import jax.numpy as jnp
 
-        >>> 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)
-        (5, 3)
-
-    Another example with a more complex model::
-
-        >>> class LinearEnsemble(bst.nn.Module):
-        ...     def __init__(self, n: int):
-        ...         super().__init__()
-        ...         self.n = n
-        ...         self.w = bst.ParamState(bst.random.random((n, 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)
-        (5, 3)
+        >>> 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))
 
-    To control how different types of states are vectorized, ``StateAxes``
-    can be passed to ``in_axes`` and ``out_axes`` specifying the axes to be
-    applied to each substate given a filter. The following example shows how to
-    share the parameters between the ensemble members which keeping different
-    batch statistics and dropout random state::
+        >>>     def __call__(self, *args, **kwargs):
+        >>>         self.c.value = self.a.value * self.b.value
+        >>>         return self.c.value + 1.
 
-        >>> class Foo(bst.nn.Module):
-        ...     def __init__(self):
-        ...         super().__init__()
-        ...         self.a = bst.ParamState(jnp.arange(4))
-        ...         self.b = bst.ShortTermState(jnp.arange(4))
+        >>> model = Model()
 
-        >>> 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
-
-        >>> model = Foo()
-        >>> y = mul(model)
-        >>> print(y.shape)
-        (4, 4)
+        >>> r = bst.augment.vmap(
+        >>>     model,
+        >>>     in_states=model.states(bst.ShortTermState),
+        >>>     out_states=model.c
+        >>> )()
 
     Args:
         fn: Function to be mapped over additional axes.
         in_axes: An integer, None, or sequence of values specifying which input
           array axes to map over.
-
-          If each positional argument to ``fun`` is an array, then ``in_axes`` can
-          be an integer, a None, or a tuple of integers and Nones with length equal
-          to the number of positional arguments to ``fun``. An integer or ``None``
-          indicates which array axis to map over for all arguments (with ``None``
-          indicating not to map any axis), and a tuple indicates which axis to map
-          for each corresponding positional argument. Axis integers must be in the
-          range ``[-ndim, ndim)`` for each array, where ``ndim`` is the number of
-          dimensions (axes) of the corresponding input array.
-
-          If the positional arguments to ``fun`` are container (pytree) types, ``in_axes``
-          must be a sequence with length equal to the number of positional arguments to
-          ``fun``, and for each argument the corresponding element of ``in_axes`` can
-          be a container with a matching pytree structure specifying the mapping of its
-          container elements. In other words, ``in_axes`` must be a container tree prefix
-          of the positional argument tuple passed to ``fun``. See this link for more detail:
-          https://jax.readthedocs.io/en/latest/pytrees.html#applying-optional-parameters-to-pytrees
-
-          Either ``axis_size`` must be provided explicitly, or at least one
-          positional argument must have ``in_axes`` not None. The sizes of the
-          mapped input axes for all mapped positional arguments must all be equal.
-
-          Arguments passed as keywords are always mapped over their leading axis
-          (i.e. axis index 0).
-
-          See below for examples.
-
         out_axes: An integer, None, or (nested) standard Python container
           (tuple/list/dict) thereof indicating where the mapped axis should appear
-          in the output. All outputs with a mapped axis must have a non-None
-          ``out_axes`` specification. Axis integers must be in the range ``[-ndim,
-          ndim)`` for each output array, where ``ndim`` is the number of dimensions
-          (axes) of the array returned by the :func:`vmap`-ed function, which is one
-          more than the number of dimensions (axes) of the corresponding array
-          returned by ``fun``.
+          in the output.
         axis_name: Optional, a hashable Python object used to identify the mapped
           axis so that parallel collectives can be applied.
         axis_size: Optional, an integer indicating the size of the axis to be
@@ -296,6 +233,8 @@ def vmap(
             generators to be used in the mapped function. These random number
             generators are restored their random key after the mapped function is
             executed.
+        in_states: Optional, the :class:`State` objects to be mapped over in the inputs.
+        out_states: Optional, the :class:`State` objects to be mapped over in the outputs.
 
     Returns:
         Batched/vectorized version of ``fun`` with arguments that correspond to
@@ -304,23 +243,26 @@ def vmap(
         with extra array axes at positions indicated by ``out_axes``.
 
     """
+
     if isinstance(fn, Missing):
         return functools.partial(
-            vmap,
+            _vmap_transform,
             in_axes=in_axes,
             out_axes=out_axes,
+            in_states=in_states,
+            out_states=out_states,
             axis_name=axis_name,
             axis_size=axis_size,
             spmd_axis_name=spmd_axis_name,
             rngs=rngs,
         )  # type: ignore[return-value]
 
-    return _map_transform(
-        'vmap',  # ctxtag
-        jax.vmap,
+    return _vmap_transform(
         fn,
         in_axes=in_axes,
         out_axes=out_axes,
+        in_states=in_states,
+        out_states=out_states,
         axis_name=axis_name,
         axis_size=axis_size,
         spmd_axis_name=spmd_axis_name,
@@ -363,65 +305,6 @@ def pmap(
 
     More information please see `jax.vmap <https://jax.readthedocs.io/en/latest/_autosummary/jax.vmap.html>`__.
 
-    If there are 4 XLA devices available, the following example will execute
-    the function in parallel on each device::
-
-
-        >>> import brainstate as bst
-        >>> import jax.numpy as jnp
-
-        >>> model = bst.nn.Linear(2, 3)
-        >>> x = jnp.ones((4, 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)
-        (4, 3)
-
-    Another example with a more complex model::
-
-        >>> class LinearEnsemble(bst.nn.Module):
-        ...     def __init__(self, n: int):
-        ...         super().__init__()
-        ...         self.n = n
-        ...         self.w = bst.ParamState(bst.random.random((n, 2, 3)))
-
-        >>> model = LinearEnsemble(4)
-        >>> 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)
-        (4, 3)
-
-    To control how different types of states are vectorized, ``StateAxes``
-    can be passed to ``in_axes`` and ``out_axes`` specifying the axes to be
-    applied to each substate given a filter. The following example shows how to
-    share the parameters between the ensemble members which keeping different
-    batch statistics and dropout random state::
-
-        >>> 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
-
-        >>> model = Foo()
-        >>> y = mul(model)
-        >>> print(y.shape)
-        (4, 4)
-
 
     Args:
       fn: Function to be mapped over argument axes. Its arguments and return
@@ -508,18 +391,123 @@ def pmap(
             rngs=rngs,
         )  # type: ignore[return-value]
 
-    return _map_transform(
-        'pmap',  # ctxtag
-        jax.pmap,
-        fn,
-        in_axes=in_axes,
-        out_axes=out_axes,
-        axis_name=axis_name,
-        static_broadcasted_argnums=static_broadcasted_argnums,
-        devices=devices,
-        backend=backend,
-        axis_size=axis_size,
-        donate_argnums=donate_argnums,
-        global_arg_shapes=global_arg_shapes,
-        rngs=rngs,
+    return restore_rngs(
+        jax.pmap(
+            fn,
+            in_axes=in_axes,
+            out_axes=out_axes,
+            axis_name=axis_name,
+            static_broadcasted_argnums=static_broadcasted_argnums,
+            devices=devices,
+            backend=backend,
+            axis_size=axis_size,
+            donate_argnums=donate_argnums,
+            global_arg_shapes=global_arg_shapes,
+        ),
+        rngs=rngs
     )
+
+
+def _batch_and_remainder(x, batch_size: int):
+    leaves, tree_def = jax.tree.flatten(x)
+
+    scan_leaves = []
+    remainder_leaves = []
+
+    length = None
+    for leaf in leaves:
+        if length is None:
+            length = leaf.shape[0]
+        if length != leaf.shape[0]:
+            raise ValueError(f"All inputs must have the same length. Got {length} and {leaf.shape[0]}.")
+
+    num_batches, num_remainder = divmod(length, batch_size)
+    for leaf in leaves:
+        total_batch_elems = num_batches * batch_size
+        scan_leaves.append(leaf[:total_batch_elems].reshape(num_batches, batch_size, *leaf.shape[1:]))
+        if num_remainder:
+            remainder_leaves.append(leaf[total_batch_elems:])
+
+    scan_tree = tree_def.unflatten(scan_leaves)
+    if num_remainder:
+        remainder_tree = tree_def.unflatten(remainder_leaves)
+        return scan_tree, remainder_tree
+    else:
+        return scan_tree, None
+
+
+def map(
+    f,
+    *xs,
+    batch_size: int | None = None,
+):
+    """
+    Map a function over leading array axes.
+
+    Like Python's builtin map, except inputs and outputs are in the form of
+    stacked arrays. Consider using the :func:`~jax.vmap` transform instead, unless you
+    need to apply a function element by element for reduced memory usage or
+    heterogeneous computation with other control flow primitives.
+
+    When ``xs`` is an array type, the semantics of :func:`~map` are given by this
+    Python implementation::
+
+        def map(f, *xs):
+            return np.stack([f(*x) for x in xs])
+
+    Like :func:`~scan`, :func:`~map` is implemented in terms of JAX primitives so
+    many of the same advantages over a Python loop apply: ``xs`` may be an
+    arbitrary nested pytree type, and the mapped computation is compiled only
+    once.
+
+    If ``batch_size`` is provided, the computation is executed in batches of that size
+    and parallelized using :func:`~jax.vmap`. This can be used as either a more performant
+    version of ``map`` or as a memory-efficient version of ``vmap``. If the axis is not
+    divisible by the batch size, the remainder is processed in a separate ``vmap`` and
+    concatenated to the result.
+
+        >>> import jax.numpy as jnp
+        >>> x = jnp.ones((10, 3, 4))
+        >>> def f(x):
+        ...   print('inner shape:', x.shape)
+        ...   return x + 1
+        >>> y = map(f, x, batch_size=3)
+        inner shape: (3, 4)
+        inner shape: (3, 4)
+        >>> y.shape
+        (10, 3, 4)
+
+    In the example above, "inner shape" is printed twice, once while tracing the batched
+    computation and once while tracing the remainder computation.
+
+    Args:
+        f: a Python function to apply element-wise over the first axis or axes of
+            ``xs``.
+        xs: values over which to map along the leading axis.
+        batch_size: (optional) integer specifying the size of the batch for each step to execute
+            in parallel.
+
+    Returns:
+        Mapped values.
+    """
+    if batch_size is not None:
+        scan_xs, remainder_xs = _batch_and_remainder(xs, batch_size)
+        g = lambda _, x: ((), vmap(f)(*x))
+        _, scan_ys = scan(g, (), scan_xs)
+        if remainder_xs is None:
+            ys = jax.tree.map(lambda x: flatten_(x), scan_ys)
+        else:
+            remainder_ys = vmap(f)(*remainder_xs)
+            ys = jax.tree.map(
+                lambda x, y: jax.lax.concatenate([flatten_(x), y], dimension=0),
+                scan_ys,
+                remainder_ys,
+            )
+    else:
+        g = lambda _, x: ((), f(*x))
+        _, ys = scan(g, (), xs)
+    return ys
+
+
+def flatten_(x):
+    return x.reshape(-1, *x.shape[2:])