brainstate.transform.vmap

Contents

brainstate.transform.vmap#

brainstate.transform.vmap(fn=<brainstate.typing.Missing object>, *, in_axes=0, out_axes=0, axis_name=None, axis_size=None, spmd_axis_name=None, in_states=None, out_states=None)[source]#

Vectorize a callable while preserving BrainState state semantics.

This is the declaration-based vectorization API: states that participate in the mapped axis are declared explicitly through in_states and out_states (by State instance). It is implemented as a thin shim over the shared mapping engine that also powers vmap2(); the declared states are converted to identity selectors internally.

Compared with vmap2(), this entry point keeps the historical contract: a state written with a batched value but not declared in out_states raises a BatchAxisError (rather than being inferred automatically), and keyword arguments are not supported.

Parameters:

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

  • in_axes (int | None | Sequence[Any]) – Mapped-axis alignment per positional argument, following jax.vmap() semantics. None marks an argument as broadcast.

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

  • 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/states when omitted.

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

  • in_states (Dict[int, Dict] | Any | None) – States batched on input, declared by instance. A dict maps axis identifiers to states; a bare state (or iterable of states) is shorthand for {0: ...}.

  • out_states (Dict[int, Dict] | Any | None) – States whose writes are scattered back along the mapped axis, with the same declaration semantics as in_states.

Returns:

The vectorized function if fn is supplied, otherwise a decorator.

Return type:

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

Raises:

  • BatchAxisError – If a state is written with a batched value but not declared in out_states.

  • NotImplementedError – If keyword arguments are passed to the vectorized function.

See also

brainstate.transform.vmap2

Filter/predicate-based vectorization with automatic output-axis inference.

brainstate.transform.vmap_new_states

Vectorize states created inside the function.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> counter = brainstate.ShortTermState(jnp.zeros(3))
>>>
>>> @brainstate.transform.vmap(
...     in_axes=0,
...     in_states=counter,
...     out_states=counter,
... )
... 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)
Contents