StatefulMapping

class brainstate.transform.StatefulMapping(fun, in_axes=0, out_axes=0, state_in_axes=None, state_out_axes=None, unexpected_out_state_mapping='auto', static_argnums=(), static_argnames=(), axis_env=None, return_only_write=True, axis_size=None, axis_name=None, name=None, mapping_fn=<function vmap>, mapping_kwargs=None)

State-aware mapping wrapper built on the shared brainstate.transform mapping engine.

StatefulMapping augments a JAX mapping primitive (jax.vmap or jax.pmap) with awareness of State instances. It tracks reads and writes across the mapped axis, splits random states so each lane is seeded independently, scatters batched writes back to the right axis, and restores the global RNG so randomness is consumed exactly once per call. It is normally constructed by brainstate.transform.vmap2() or brainstate.transform.pmap2(), but can be used directly for custom mapping primitives.

Parameters:

• fun (Callable) – Callable to wrap. May read and write closed-over State objects.

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

• out_axes (int | Tuple[int, ...] | None) – Placement of the mapped axis in the return value.

• 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]]) – Which states participate as batched inputs. A dict maps axis identifiers to selectors; a bare selector is shorthand for {0: selector}. Selectors may be brainstate.util.filter filters, a single State instance, or an iterable of instances.

• 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]]) – Which written states are scattered back along the mapped axis, with the same selector semantics as state_in_axes.

• unexpected_out_state_mapping (str) –

  Policy for states written with a batched value but not covered by state_out_axes. 'auto' scatters them at their detected axis; 'raise' raises a BatchAxisError; 'warn' scatters them with a warning; and 'ignore' scatters them silently.

  Under 'auto', an undeclared state that is read and written (read-modify-write) is handled specially: on each call, if its current leading size along the detected axis already equals the mapped size it is fed per lane (a per-lane RMW buffer) instead of being broadcast and scattered. This decision is re-made every call from the live value, so a cached (warm) call behaves identically to a fresh (cold) one. Because the choice rests on a leading-size match, a state whose leading dimension coincidentally equals the mapped size is treated as per-lane even if you meant it as a shared, broadcast value – declare it explicitly via state_in_axes/state_out_axes to remove the ambiguity. When such an RMW state’s leading size does not match the mapped size it is scattered (gaining a new leading axis) and a one-time UserWarning is emitted.

• static_argnums (int | Iterable[int]) – Positional arguments treated as compile-time constants. The argument is closed over (as in jax.jit()) and is neither traced nor mapped, so its in_axes entry, if any, is ignored.

• static_argnames (str | Iterable[str]) – Keyword arguments treated as compile-time constants for caching.

• axis_env (Sequence[tuple[Hashable, int]] | None) – Retained for backward compatibility; mapping primitives manage the axis environment via axis_name.

• return_only_write (bool) – Retained for backward compatibility.

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

• axis_name (Hashable | None) – Name of the mapped axis for collective primitives.

• name (str | None) – Diagnostic identifier.

• mapping_fn (Callable) – Mapping primitive accepting in_axes/out_axes/axis_size/ axis_name.

• mapping_kwargs (Dict) – Extra keyword arguments forwarded to mapping_fn.

origin_fun

The wrapped callable.

Type:

callable

in_axes, out_axes

Argument/return mapping specification.

Type:

int, tuple, or None

state_in_axes, state_out_axes

Normalized {axis: predicate} state selectors.

Type:

dict

axis_size

Mapped axis size, if explicitly provided.

Type:

int or None

axis_name

Axis identifier forwarded to collectives.

Type:

hashable or None

mapping_fn

The underlying mapping primitive.

Type:

callable

Notes

Random states are always split along the mapped axis and restored afterwards; this cannot be disabled. Plans (state groupings, batch sizes) are cached per abstract argument signature so repeated calls with the same structure avoid re-tracing.

Examples

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