Jax does actually already support a persistent compilation cache for TPU, and support for caching GPU compiles is being worked on currently.