LRU Cache and JAX

[SNMS95]

To my knowledge, Jax's JIT compilation of a function is stored using lru's cache.
(#8458)
I understand that subsequent calls to a JITTed function would then do a look up in this cache and run that version instead of the original. But if we run it with the same inputs, the computation is not cached right?

But putting a lru_cache decorator fails

import jax.numpy as jnp
from functools import lru_cache
import jax
import time

def perform_heavy_computation(x):
    # Simulate a heavy computation
    return jnp.sum(x ** 2)

x = jax.random.normal(key=jax.random.PRNGKey(0), shape=(10000, 10000))

# Time normal computation
start_time = time.time()
result = perform_heavy_computation(x).block_until_ready()
normal_time = time.time() - start_time
print(f"Normal computation time: {normal_time:.6f} seconds")

# Try with JIT compilation
jit_func = jax.jit(perform_heavy_computation)
# run once to compile
_ = jit_func(x).block_until_ready()
start_time = time.time()
result_jit = jit_func(x).block_until_ready()
jit_time = time.time() - start_time
print(f"JIT compilation time: {jit_time:.6f} seconds")


# With LRU cache
cached_func = lru_cache(maxsize=1)(perform_heavy_computation)
# run once to cache
_ = cached_func(x).block_until_ready()
start_time = time.time()
result_cached = cached_func(x).block_until_ready()
cached_time = time.time() - start_time
print(f"Cached computation time: {cached_time:.6f} seconds")

[jakevdp]

Your code fails because lru_cache only works with hashable arguments, and JAX arrays are not hashable.

JAX’s JIT does employ an LRU cache, but not over the function arguments themselves. Rather the cache is over the static properties of array arguments (e.g. shape, dtype, etc.) so if you call the function a second time with arrays of the same shape and type (even if they don’t have the same contents) you’ll execute without recompilation.

I jope that’s clear!