[Python] Jax persistent cache breaks causes miscompilation?

When I comment the below line ("THIS LINE"), I get correct results (loss decreases), and when I don't, I get broken results (constant loss every iteration):

import jax
import numpy as np
from jax import numpy as jnp
import optax
from tqdm import tqdm
from pathlib import Path
import shutil

cache_path = "/tmp/jax_cache"
if Path(cache_path).exists():
shutil.rmtree(cache_path)
jax.clear_caches()
jax.config.update("jax_compilation_cache_dir", cache_path)
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0) # THIS LINE

print('initializing...')
params = jnp.array(np.random.normal(size=[10, 11]))

def loss(params):
return jnp.sum(params ** 4) / 1000.0

print('training...')
opt = optax.lbfgs(0.03)

@jax.jit
def do_update(params, opt_state):
loss_value, params_grad = jax.value_and_grad(loss)(params)
updates, opt_state = opt.update(
params_grad,
opt_state,
params=params,
value=loss_value,
grad=params_grad,
value_fn=loss,
)
params = optax.apply_updates(params, updates)
return params, opt_state

opt_state = opt.init(params)

def maybe_log(i):
if i % 100 == 0:
loss_float = loss(params).item()
tqdm.write(f'{i} {loss_float}')

for i in tqdm(range(1000)):
maybe_log(i)
params, opt_state = do_update(params, opt_state)


Why is this happening? Is it a bug in jax or a problem in my code?

My understanding is that this line causes jax to put do_update in the persistent cache, so it's surprising to me that this change's the code's behavior. I've believe I've even cleared all caches at the start of this code.

• jax: 0.5.3
• hardware: aarch64 CPU
• python: 3.13.0

Continue reading...