brainstate.transform.pmap2_new_states

brainstate.transform.pmap2_new_states#

brainstate.transform.pmap2_new_states(module, init_kwargs, state_tag=None, axis_size=None, state_out_axes=None, axis_name=None)#

Initialize and parallelize newly created states across devices.

Wraps module.init_all_states(**init_kwargs) in a pmap2()-style transform, executes it across axis_size devices, and restores the device-distributed values back onto the freshly created state objects. Random states are split per device.

Note

init_all_states runs twice (a probe pass plus the mapped pass) and must be idempotent; see _map_new_states() for details.

Parameters:

  • module (Module) – Module whose init_all_states creates the states to parallelize.

  • init_kwargs (Dict) – Keyword arguments forwarded to module.init_all_states.

  • state_tag (str) – Tag applied to the newly created states.

  • axis_size (int) – Number of devices to map over. Defaults to jax.local_device_count().

  • state_out_axes (Dict[int, Filter]) – Output-axis selectors. None (default) shards every state on axis 0 except NonBatchState, which is replicated.

  • axis_name (AxisName | None) – Mapped-axis name for collective primitives.

Returns:

Mapping of axis identifier to the lists of parallelized states.

Return type:

Dict

See also

brainstate.transform.pmap2

Parallel mapping across devices.

brainstate.transform.vmap2_new_states

Single-device variant.

jax.pmap

Underlying JAX parallel mapping primitive.

Examples

>>> import brainstate
>>> import jax
>>> import jax.numpy as jnp
>>>
>>> class ParallelCounter(brainstate.nn.Module):
...     def init_state(self):
...         self.count = brainstate.ShortTermState(jnp.zeros(()))
>>>
>>> module = ParallelCounter()
>>> _ = brainstate.transform.pmap2_new_states(
...     module, init_kwargs={}, axis_size=jax.local_device_count()
... )
>>> module.count.value.shape
(jax.local_device_count(),)