quantax.utils.chunk_map#

quantax.utils.chunk_map(f: Callable, in_axes: tuple | int | None = 0, out_axes: tuple | int | None = 0, chunk_size: int | None = None, use_scan: bool = False, shard_batch: bool = False) → Callable#

Convert a vmapped function to a function with chunked batches and parallel computation on all available machines. The arguments will be unchanged if the batch size on each machine is smaller than the chunk size, but it will be padded with 0 if the batch size is larger than the chunk size and not a multiple of chunk size.

Parameters:

  • f – The function to be converted. The arguments of f are assumed to be sharded.

  • in_axes – The vmapped axes of f which are to be chunked.

  • out_axes – The vmapped axes of outputs.

  • chunk_size – The chunk size on each machine.

  • use_scan – Whether to use jax.lax.scan in chunked function apply. The compilation will be accelerated if scan is used, but the function must be jittable.

  • shard_batch – Whether to run f inside jax.shard_map so each device chunks and processes only its own batch shard (a device-local lax.scan inside one shard_map), instead of GSPMD partitioning a global batch. Needed only for per-sample reverse-mode ops where GSPMD inserts cross-device all-gathers – the symmetry-vmapped Jacobian’s conv weight-gradient – so it defaults to False. Must stay False when f calls another chunk_map-wrapped function, because shard_map cannot be nested over the same mesh axes.