quantax.utils.use_portable_compilation_cache

quantax.utils.use_portable_compilation_cache() → None

Make persistent-compilation-cache keys portable across device topologies of the same hardware, so executables written by precompile under a compile-only mesh are found by real runs.

Stock jax hashes the full device-topology fingerprint into every cache key. The fingerprint of a compile-only topology differs from that of the same devices in a real multi-node run (host layout metadata enters the fingerprint) even though the compiled executable is identical, so the real run would never find the precompiled entries. This replaces the topology component of the key with the device model(s) and the device count. All other components – the computation, the compile options (including partition count and device assignment), the jaxlib version, the CUDA version, and the XLA flags – are hashed as usual.

Call this in every script that shares the cache across jobs: the precompiling job applies it automatically through make_precompile_mesh(), and the target run must call it explicitly before the first compilation.