stringjax_tools.cache.configure_compilation_cache#
- configure_compilation_cache(cache_dir=None, *, max_size_bytes=None, min_compile_time_secs=None, min_entry_size_bytes=None, enable_xla_caches=None, explain_cache_misses=None)#
Configure JAX’s persistent compilation cache explicitly.
This helper has no import-time side effects. Call it near the top of a script or notebook before any JAX compilation happens.
- Parameters:
cache_dir (
str|Path|None) – Directory for the persistent compilation cache. If supplied, the directory is created if needed. IfNone, the cache directory is not changed.max_size_bytes (
int|None) – Maximum persistent cache size in bytes. JAX interprets0as disabling the persistent cache and-1as no size limit.min_compile_time_secs (
float|None) – Minimum compilation time required before an entry is written to the persistent cache.min_entry_size_bytes (
int|None) – Minimum entry size required before an entry is written to the persistent cache.enable_xla_caches (
str|None) – Value forwarded tojax_persistent_cache_enable_xla_caches.explain_cache_misses (
bool|None) – IfTrue, ask JAX to explain persistent-cache misses.
- Returns:
Configuration names and values that were applied.
- Return type:
dict[str,Any]