brainstate.transform.map

brainstate.transform.map(f, *xs, batch_size=None)

Apply a function over the leading axis of one or more pytrees.

Compared with jax.vmap(), this helper executes sequentially by default (via brainstate.transform.scan()), which keeps peak memory low. When batch_size is given, full batches are processed with vmap2() and any remainder is handled separately, trading memory for throughput.

Parameters:

• f (Callable) – Function applied across the leading dimension. Its return value must be a pytree whose leaves can be stacked along axis 0.

• *xs (Any) – Positional pytrees sharing the same leading length.

• batch_size (int | None) – Size of vectorised blocks. When given, map processes full batches with vmap2() and then handles any remainder. The batched path treats f as stateless (a pure throughput optimization); a function that writes a non-random State is rejected (see Raises). Drawing randomness is allowed – random states thread correctly through the batched path. The sequential path (no batch_size) handles state writes normally.

Returns:

A pytree matching the structure of f’s outputs, stacked along the leading dimension.

Return type:

Any

Raises:

ValueError – If the inputs do not share the same leading length, if batch_size is not a positive integer, or if batch_size is given and f writes a non-random State (use sequential map or vmap2() instead). Random-number draws are permitted.

See also

brainstate.transform.vmap2

Vectorised mapping with automatic batching.

brainstate.transform.scan

Primitive used for the sequential path.

Examples

>>> import jax.numpy as jnp
>>> from brainstate.transform import map
>>>
>>> xs = jnp.arange(6).reshape(6, 1)
>>>
>>> def normalize(row):
...     return row / (1.0 + jnp.linalg.norm(row))
>>>
>>> map(normalize, xs, batch_size=2).shape
(6, 1)