brainstate.transform.vmap2

brainstate.transform.vmap2(fn=<brainstate.typing.Missing object>, *, in_axes=0, out_axes=0, axis_name=None, axis_size=None, spmd_axis_name=None, state_in_axes=None, state_out_axes=None, unexpected_out_state_mapping='auto')

Vectorize a callable while preserving BrainState state semantics.

This mirrors jax.vmap() but routes execution through StatefulMapping so reads and writes to State instances (including random states) are tracked across the mapped axis. The returned object can be used directly or as a decorator when fn is omitted.

Parameters:

• fn (TypeVar(F, bound= Callable) | Missing) – Function to vectorise. If omitted, the function acts as a decorator.

• in_axes (int | Sequence[Any] | None) – Mapping specification for positional arguments.

• out_axes (Any) – Placement of the mapped axis in the result.

• axis_name (Hashable | None) – Name for the mapped axis so collective primitives can target it.

• axis_size (int | None) – Explicit mapped axis size. Inferred from arguments when omitted.

• spmd_axis_name (Hashable | Tuple[Hashable, ...] | None) – Axis labels for nested SPMD transforms.

• state_in_axes (Dict[Hashable, type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter], ...] | List[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]]] | type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter], ...] | List[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]]) – Selectors for states batched on input. A bare selector means {0: ...}.

• state_out_axes (Dict[Hashable, type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter], ...] | List[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]]] | type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter], ...] | List[type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]]) – Selectors for written states scattered back along the mapped axis.

• unexpected_out_state_mapping (str) – Policy for states written with a batched value but not declared in state_out_axes. The default 'auto' infers the output axis from the detected batch dimension. For an undeclared read-modify-write state, 'auto' additionally feeds it per lane when its current leading size matches the mapped size (re-checked every call, so warm and cold calls agree), and warns once when it does not (the state is scattered, gaining a new axis). See StatefulMapping for the full description and how to disambiguate a coincidental size match with state_in_axes/state_out_axes.

Returns:

A StatefulMapping if fn is supplied, otherwise a decorator.

Return type:

StatefulMapping | Callable[[TypeVar(F, bound= Callable)], StatefulMapping]

Raises:

• ValueError – If axis sizes are inconsistent or cannot be inferred.

• BatchAxisError – If a state write violates state_out_axes under the 'raise' policy.

See also

brainstate.transform.StatefulMapping

Underlying state-aware mapping helper.

brainstate.transform.pmap2

Multi-device parallel variant.

brainstate.transform.vmap

Declaration-based vectorisation (now a shim).

Notes

Keyword arguments passed when calling the wrapped function are mapped over axis 0, matching jax.vmap() (they are not broadcast). To keep a keyword argument broadcast as a compile-time constant, construct a StatefulMapping directly and list it in static_argnames.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>> from brainstate.util.filter import OfType
>>>
>>> counter = brainstate.ShortTermState(jnp.zeros(3))
>>>
>>> @brainstate.transform.vmap2(
...     in_axes=0,
...     out_axes=0,
...     state_in_axes={0: OfType(brainstate.ShortTermState)},
...     state_out_axes={0: OfType(brainstate.ShortTermState)},
... )
... def accumulate(x):
...     counter.value = counter.value + x
...     return counter.value
>>>
>>> accumulate(jnp.asarray([1., 2., 3.]))
Array([1., 2., 3.], dtype=float32)
>>> counter.value
Array([1., 2., 3.], dtype=float32)