quantax.utils.make_precompile_mesh#

quantax.utils.make_precompile_mesh(num_processes: int, local_device_count: int) → Mesh#

Build a compile-only ("process", "device") mesh describing the topology of a different (typically larger) run, e.g. a multi-node job, without owning its devices. Ahead-of-time compilations staged on this mesh (see precompile) are written to the persistent compilation cache with the same keys the real run computes, so the real run loads the executables instead of compiling them.

Parameters:

  • num_processes – Number of processes of the target run (typically the number of nodes).

  • local_device_count – Number of devices per process of the target run (typically GPUs per node).

Returns:

A compile-only ("process", "device") jax.sharding.Mesh with shape (num_processes, local_device_count), mirroring make_mesh() of the target run.

Warning

This is only supported on GPU, and requires at least one visible GPU of the same model, CUDA version, jaxlib version, and XLA flags as the target run — cache keys include all of them, so in practice the precompiling job should run on one node of the same cluster with the same environment.

Note

Arrays cannot be created on a compile-only mesh; it can only be used to trace, lower, and compile. Executing a computation staged on it raises an error.