brainstate.transform.vmap_new_states

brainstate.transform.vmap_new_states#

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

Vectorize a function over the new states it creates.

Unlike vmap(), which maps over states that already exist, this transform maps over states created during the call to fun (for example, parameters allocated inside a module’s initializer). Each mapped lane creates its own copy of every new state, and random states are split per lane so randomly initialized values differ across the mapped axis. It is implemented as a single mapping pass over the shared engine helpers.

Note

fun is executed twice – an eager probe (to discover the random states it touches) plus the mapped pass – so it must be idempotent and free of un-rolled-back side effects. NonBatchState (and INIT_NO_BATCHING-tagged) states created inside fun are replicated rather than batched.

Parameters:

  • fun (Callable) – Function to vectorize. If omitted, this 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 the inputs when omitted.

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

  • state_tag (str | None) – Tag applied to the newly created states so they can be retrieved later.

  • state_to_exclude (type | str | Callable[[Tuple[Key, ...], Any], bool] | bool | EllipsisType | None | Tuple[Filter, ...] | List[Filter]) – Selector for new states that should be left untouched (not vectorized).

  • in_states (Dict[int, Dict] | Any | None) – Not supported by this transform. vmap_new_states maps over states created inside fun, so there are no pre-existing states to declare. Passing either raises ValueError; use vmap() to vectorize over pre-existing states.

  • out_states (Dict[int, Dict] | Any | None) – Not supported by this transform. vmap_new_states maps over states created inside fun, so there are no pre-existing states to declare. Passing either raises ValueError; use vmap() to vectorize over pre-existing states.

Returns:

A vectorized version of fun that handles new-state creation, or a decorator if fun is omitted.

Return type:

Callable

Raises:

  • ValueError – If axis_size is provided but not a positive integer, or if in_states / out_states is provided (use vmap() for pre-existing states).

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

See also

brainstate.transform.vmap

Vectorize over pre-existing states.

brainstate.transform.vmap2_new_states

Module-oriented new-state mapping.

Examples

>>> import brainstate
>>> import jax.numpy as jnp
>>>
>>> @brainstate.transform.vmap_new_states(in_axes=0, axis_size=4)
... def build(x):
...     scratch = brainstate.ShortTermState(jnp.zeros(()))
...     scratch.value = scratch.value + x
...     return scratch.value
>>>
>>> build(jnp.arange(4.))
Array([0., 1., 2., 3.], dtype=float32)