jax.experimental.compilation_cache module

JAX disk compilation cache.

API

jax.experimental.compilation_cache.compilation_cache.is_initialized()

Deprecated.

Return whether the cache is enabled. Initialization can be deferred, so initialized status is not checked. The name is retained for backwards compatibility.

Return type:

bool

jax.experimental.compilation_cache.compilation_cache.initialize_cache(path)

This API is deprecated; use set_cache_dir instead.

Set the path. To take effect, should be called prior to any calls to get_executable_and_time() and put_executable_and_time().

Return type:

None

jax.experimental.compilation_cache.compilation_cache.set_cache_dir(path)

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.

Return type:

None

jax.experimental.compilation_cache.compilation_cache.reset_cache()

Get back to pristine, uninitialized state.

Return type:

None