brainstate.transform.pmap2

brainstate.transform.pmap2(fn=<brainstate.typing.Missing object>, axis_name=None, *, in_axes=0, out_axes=0, static_broadcasted_argnums=(), devices=None, backend=None, axis_size=None, donate_argnums=(), global_arg_shapes=None, state_in_axes=None, state_out_axes=None, unexpected_out_state_mapping='auto')

Parallel-map a callable across devices with state-aware semantics.

This mirrors jax.pmap() but integrates with StatefulMapping so State objects (including random states) are replicated, split, and restored correctly on every device. When fn is omitted a decorator is returned.

Parameters:

• fn (Callable[[NestedDict, ...], Any] | Missing) – Function to execute in SPMD style. If omitted, a decorator is returned.

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

• in_axes (Any) – Axis mapping for positional arguments.

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

• static_broadcasted_argnums (int | Iterable[int]) – Not supported. Positional-index static broadcasting cannot be honored because the state-aware wrapper bundles your arguments before calling jax.pmap(). A non-empty value raises NotImplementedError; use static_argnames or jax.pmap() directly instead.

• devices (Sequence[Device] | None) – Explicit device list to map over.

• backend (str | None) – Backend identifier ('cpu', 'gpu', or 'tpu').

• axis_size (int | None) – Size of the mapped axis. Defaults to the local device count.

• donate_argnums (int | Iterable[int]) – Not supported. Buffer donation by positional index cannot be honored for the same reason as static_broadcasted_argnums; a non-empty value raises NotImplementedError. Call jax.pmap() directly if you need donation.

• global_arg_shapes (Tuple[Tuple[int, ...], ...] | None) – Shapes for globally distributed arguments.

• 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 treated as device-mapped inputs.

• 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 state writes scattered back across devices.

• unexpected_out_state_mapping (str) – Policy for state writes not covered by state_out_axes.

Returns:

A StatefulMapping executing fn over devices, or a decorator.

Return type:

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

See also

jax.pmap

Underlying JAX primitive.

brainstate.transform.vmap2

Single-host vectorisation with the same semantics.