Compilation-cache setup

configure_compilation_cache is an opt-in helper for JAX’s persistent compilation cache. It has no import-time side effects. Call it near the top of a script or notebook, before the first expensive JAX compilation in that Python process.

from pathlib import Path
import tempfile

import jax
import jax.numpy as jnp

from stringjax_tools.cache import configure_compilation_cache

The helper simply applies the requested JAX config values and returns the keys that were updated. The temporary directory below keeps the example isolated.

cache_dir = Path(tempfile.mkdtemp(prefix='stringjax_tools_cache_'))
updates = configure_compilation_cache(
    cache_dir=cache_dir,
    max_size_bytes=-1,
    min_compile_time_secs=0.0,
    min_entry_size_bytes=0,
)

for key, value in updates.items():
    print(f'{key}: {value}')

After configuration, ordinary JAX compilation proceeds as usual. Whether a particular compiled executable is persisted depends on JAX and backend cache rules.

@jax.jit
def norm2(x):
    return jnp.sum(x * x)

print(norm2(jnp.arange(5.0)))