jax.experimental.compilation_cache module#
JAX disk compilation cache.
API#
- jax.experimental.compilation_cache.compilation_cache.set_cache_dir(path)[source]#
Sets the persistent compilation cache directory.
After calling this, jit-compiled functions are saved to path, so they do not need be recompiled if the process is restarted or otherwise run again. This also tells Jax where to look for compiled functions before compiling.
For more information, see the persistent compilation cache guide.
Warning
The compilation cache is considered trusted. Do not share a compilation cache with users you do not trust. For example, if you put the compilation cache in a directory to which others may write, those users can trigger your JAX process to run arbitrary code. Sharing a compilation cache is equivalent to allowing anyone who can write to the cache directory to run code on your machine.
- Return type:
None
- jax.experimental.compilation_cache.compilation_cache.reset_cache()[source]#
Get back to pristine, uninitialized state.
For more information, see the persistent compilation cache guide.
- Return type:
None