adv_jax_math.batch_vmap

adv_jax_math.batch_vmap(f, /, in_axes=0, *, batch_size=None, reduction=None, chunk_reduction=<function identity>, shard=False, mesh=None, **kwargs)[source]

Behaves like vmap but uses scan to chunk the computations in smaller chunks.

Warning

See also

batch_map

If the function supports native vectorization, use batch_map instead for the reasons discussed the docstring.

Parameters:

  • f (callable) – The function to be vectorised.

  • in_axes (int or None) – The axes that should be scanned along. Only supports 0 or None.

  • batch_size (int or None) – Size to split computation into chunks. If no chunking should be done or the chunk 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.

  • shard (bool) – Whether to shard mapped input data 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 mapped 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:

f (callable) – A vectorised and chunked function.