quantax.utils.shmap

quantax.utils.shmap(f: Callable, in_axes: tuple | int | None, out_axes: tuple | int | None, mesh: Mesh | None = None) → Callable

f -> shard_map(f), sharded along the first dimension

Parameters:

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

• in_axes – The sharded axes of f input arguments.

• out_axes – The sharded axes of f outputs.

• mesh – The device mesh to shard over, default to make_mesh spanning all devices of the current run. Pass a compile-only mesh from make_precompile_mesh to stage the computation for a different topology without running on it.