Skip to content

← JAX Performance & Distributed step 23 of 25

Medium Primitives

compilation_cache.set_cache_dir

Jump to the task

Keep compiled code between runs

Ordinary JIT reuse keeps compiled functions available within a Python process. A persistent cache saves compiled programs to disk so a later process can reuse compatible entries. It stores executable code, not answers: new input values still need to be evaluated.

Configure a private cache directory before the first compilation. jax.config.update("jax_compilation_cache_dir", path) and compilation_cache.set_cache_dir(path) are supported ways to set it. Backend initialization is not the same event as compilation.

Why a tiny example might not be saved

Cache writes depend on compilation-time and entry-size thresholds. The default minimum compilation time is one second, so a small sum may not qualify. For a local experiment, these settings remove both thresholds; they are not a recommendation to cache every tiny operation in an application.

from pathlib import Path
import jax
import jax.numpy as jnp

cache_dir = str(Path.home() / ".cache" / "jax-lesson")
jax.config.update("jax_compilation_cache_dir", cache_dir)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)

@jax.jit
def f(x):
    return jnp.sum(x ** 2 + 1)

result = f(jnp.array([1., 2., 3.]))
print(result.block_until_ready())  # 17

Run the script twice using the same directory and environment. A compatible stored entry can avoid backend compilation, though tracing, loading and execution still cost time. Matching depends on the program, jaxlib version, compiler flags and device configuration, among other details. A cache hit is not guaranteed merely by reusing a filename.

Ordinary repeated calls usually already benefit from in-memory reuse. To investigate disk reuse, inspect JAX cache logs in separate processes; a faster second call alone does not establish a disk-cache hit. Only use cache directories whose writers you trust: cached executables are trusted code.

Problem

Implement jit_with_warm_cache(x) using a module-level helper decorated with @jax.jit. The helper computes jnp.sum(x ** 2 + 1); the entrypoint calls that existing helper and returns its scalar JAX result.

X is a one-dimensional float32 JAX array. For [1, 2, 3]:

  1. Square each entry to get [1, 4, 9].
  2. Add one to every entry to get [2, 5, 10].
  3. Sum to get 17.

For [3, 4], the result is 27. For an empty array, it is zero. sum(x ** 2) + 1 adds one only once and is a different calculation.

Do not configure a disk directory in your submitted function. The cache experiment above is separate application setup. Despite its name, this entrypoint does not promise a warm cache. Tests check the numerical result; they cannot establish disk writes, cross-process reuse or speed.