jax.experimental.compilation_cache module

JAX disk compilation cache.

API

jax.experimental.compilation_cache.compilation_cache.is_initialized()

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)

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.reset_cache()

Get back to pristine, uninitialized state.

Return type:

None