3 ms·
The same issue exists with Jax. XLA compilation can take up quite a bit of time, especially on larger NN models. And theres no persistent compile cache, so even
by calaphos 5y ago
The same issue exists with Jax. XLA compilation can take up quite a bit of time, especially on larger NN models. And theres no persistent compile cache, so even if you don't change the jitted function you need to wait for compilation again as you restart the process.
- alevskaya 5y agoJax does actually already support a persistent compilation cache for TPU, and support for caching GPU compiles is being worked on currently.