brainstate.transform.vmap2_new_states

brainstate.transform.vmap2_new_states#

brainstate.transform.vmap2_new_states(module, init_kwargs, state_tag=None, axis_size=None, state_out_axes=None, spmd_axis_name=None)#

Initialize and vectorize newly created states within a module.

Wraps module.init_all_states(**init_kwargs) in a vmap2()-style transform, executes it axis_size times in parallel, and restores the vectorized values back onto the freshly created state objects. Random states are split per lane, so random initializers produce a distinct draw for every batch member.

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 vectorize.

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

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

  • axis_size (int) – Size of the vectorization axis. Required.

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

  • spmd_axis_name (AxisName | Tuple[AxisName, …] | None) – SPMD axis label forwarded to the underlying jax.vmap.

Returns:

Mapping of axis identifier to the lists of vectorized states.

Return type:

Dict

Raises:

ValueError – If axis_size is not provided.

See also

brainstate.transform.vmap2

Vectorize a callable with state semantics.

brainstate.transform.pmap2_new_states

Multi-device variant.

brainstate.NonBatchState

Marker for states that should be replicated.

Examples

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