stringjax_tools.cache.configure_compilation_cache

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. If None, the cache directory is not changed.

  • max_size_bytes (int | None) – Maximum persistent cache size in bytes. JAX interprets 0 as disabling the persistent cache and -1 as 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 to jax_persistent_cache_enable_xla_caches.

  • explain_cache_misses (bool | None) – If True, ask JAX to explain persistent-cache misses.

Returns:

Configuration names and values that were applied.

Return type:

dict[str, Any]