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, atleast_1chunk: bool = False, mesh: Mesh | None = None) → 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.

  • atleast_1chunk – Whether to pad the batch with 0 to the chunk size when the batch size on each machine is smaller than the chunk size, so that f is always called with batch size chunk_size on each machine. The padded outputs are truncated. Default to False, in which case f is called with the original batch.

  • mesh – The device mesh that determines the device count for chunking and the shard_batch sharding, default to make_mesh spanning all devices of the current run (resolved at trace time). Pass a compile-only mesh from make_precompile_mesh to stage the computation for a different topology without running on it.