adv_jax_math.batch_map

adv_jax_math.batch_map(fun, fun_input, /, batch_size=None, *, reduction=None, chunk_reduction=<function identity>, strip_dim0=False, shard=False, mesh=None, **kwargs)

Compute chunk_reduction(fun(fun_input)) in batches.

Notes

This method does not automatically wrap fun with vmap. Unless fun is already wrapped with vmap, the leading dimension of fun_input will not be stripped before it is passed into fun. This can be inconvenient for nesting calls to batch_map, since only batching along the first axis is supported. However, the strip_dim0 flag should cover the most common case of nesting calls where batch_size is one on the outermost call.

For a fun that accepts scalar and batched inputs elementwise, batch_map(fun,inputs) is expected to match batch_vmap(fun)(inputs). Functions whose results depend on other examples in a chunk or on the chunk length do not satisfy this elementwise contract.

If fun is natively vectorized, this can be preferable to batch_vmap to reduce compilation time, avoid issues such as executing all branches of code conditioned on dynamic values, or avoid messing up the behavior of jvp’s and vjp’s under vmap, e.g. https://docs.jax.dev/en/latest/jep/2026-custom-derivatives.html #main-problem-descriptions.

Only out_axes=0 is supported.

See also

batch_vmap

If the function does not support native vectorization.

Parameters:

• fun (callable) – Vectorized function.

• fun_input (pytree) – Data to split into batches to feed to fun.

• batch_size (int or None) – Size of batches. If no batching should be done or the batch size is the full input then supply None. chunk_size is accepted as an alias when batch_size is None.

• reduction (callable or None) – Binary reduction operation. Should take two arguments and return one output, e.g. jnp.add, and must be associative because partial chunk results are combined in an implementation-dependent order.

• chunk_reduction (callable) – Chunk-wise reduction operation. Should summarize a chunk compatibly with reduction along the mapped axis, e.g. jnp.add.reduce.

• strip_dim0 (bool) – Whether to strip the leading dim of fun_input before passing it to fun; see notes. This flag only works if batch_size is one. It should be set to False if fun is wrapped in vmap. Default is False.

• shard (bool) – Whether to shard fun_input across devices before applying chunked batching. The divisible prefix is split across devices; when supplied, batch_size bounds the batches processed on each device. A local remainder is evaluated once per device, and a final global remainder is evaluated once overall. The input length need not be divisible by either the device count or batch_size. Default is False. If a non-reduced output has a global remainder, the final concatenated output is replicated across the mesh.

• mesh (jax.sharding.Mesh or None) – Optional one-dimensional mesh with an AxisType.Auto axis. Supplying a mesh selects the devices and topology used by shard=True and lets already-sharded inputs retain a compatible layout. Default is None, which creates a mesh over all available devices.

Returns:

fun_output – Returns chunk_reduction(fun(fun_input)).