quantax.utils.jit_chunk_vmap#

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

f -> jit(chunk_map(vmap(f), use_scan=True))

Parameters:

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

  • in_axes – The mapped axes of f input arguments.

  • out_axes – The mapped axes of f outputs.

  • chunk_size – The chunk size on each machine. If None, no chunking will be applied.

  • shard_batch – Forwarded to chunk_map(); whether to run the vmapped f inside jax.shard_map. Defaults to False; set True only for the per-sample Jacobian.