Strange recompilations

[botev]

I'm seeing a strange behaviour of jax.jit recompiling functions in places I don't expect. Unfortunately, I don't have a simple repro to post, but roughly what I'm seeing:

1. I jit the function used to eval a model for running evals not training.
2. Running with frequency X the evals, I see only one compilation on the first time I run the evals (note the datasets are running deterministic pipelines)
3. I increase the frequency by 4, and now everytime I run evals I see a recompilation happening

[jakevdp]

Hi - thanks for the report. One debugging tool that's useful in cases like this is the JAX_EXPLAIN_CACHE_MISSES config.

Here's a quick example:

import jax
import jax.numpy as jnp
jax.config.update('jax_explain_cache_misses', True)

@jax.jit
def f(x):
    return x

x1 = f(jnp.arange(10))
x2 = f(jnp.arange(100))

This is the log output:

WARNING:2025-10-09 09:07:15,838:jax._src.pjit:1274: TRACING CACHE MISS at /Users/vanderplas/github/google/jax/tmp.py:9:5 (<module>) costing 0.075 ms because:
  never seen function:
    f id=4345667936 defined at /Users/vanderplas/github/google/jax/tmp.py:5

WARNING:2025-10-09 09:07:15,854:jax._src.pjit:1274: TRACING CACHE MISS at /Users/vanderplas/github/google/jax/tmp.py:10:5 (<module>) costing 0.170 ms because:
  for f defined at /Users/vanderplas/github/google/jax/tmp.py:5
  all previously seen cache keys are different. Closest previous key:
  * key with different input types:
      types now: x: i32[100]
        * at x, now i32[100] and before i32[10]

You can see that the first cache miss (i.e. jit compilation) happens because the function has not yet been seen by the compiler, and the second cache miss happens because x has a different shape.

Hopefully you can enable that config in your code and find the answer in the logs!

[botev]

Is there a way to enable this for internal functions as well, as I don't see this for instance for things like jnp.mean or e.g. the internal _mean which is jitted?

To maybe clarify, as I'm doing something very unorthodox, I'm trying to write a custom lowering, and atm just as a test I'm hacking and replacing pxla.lower_sharding_computation to return my own object. I would have assumed that the lowering.compile is called only on a cache miss, but I findfor instance for the jnp.mean(x) where x is a float32[100] object to get called into the compile function multiple times in my case and I can see in the tracing mechanics I can see this is being invoked via cache_miss. But then the closed_jaxpr is literally the same.

To clarify, the same Lowering object (as in id(self) is the same), returned from pxla.lower_sharding_computation is being called twice with compile, are there explanations as to why?

[jakevdp]

The explain_cache_misses flag is hard-coded to ignore internal files (otherwise there's so much logspam it's not useful). The key line is here:

jax/jax/_src/pjit.py

Lines 1269 to 1270 in fd9d71c

	if func_filename and not source_info_util.is_user_filename(func_filename):
	return

If you have a local version of JAX, you could comment-out these lines and then it would log internal cache misses as well. You may also have to comment out the "inline" check a few lines above.

[botev]

That definitely helped! However, the cache_miss is not being triggered, but the call to cache_miss is!
Let me explain, what I'm seeing, after the first call:

pjit._cpp_pjit.cache_miss is trigered and I'm trying to figure out why?
it goes to _python_pjit_helper which calls _infer_params, use_resource_env is False, so we go in _infer_params_internal. Now here _infer_params_cached doesn't hit a cache miss, but returns correctly whatever was stored from previous call, and hence we return without calling _infer_params_impl, so not cache-miss is being printed from Jax. This brings us back to pjit._cpp_pjit.cache_miss which calls _pjit_call_impl_python which should trigger util.test_event("jit_cpp_cache_miss")?. Now this goes all the way down to _pjit_lower which agains makes a cache-hit and returns the exact same Lowering from its cache and the after that on line 1782 we get the second call to .compile() on the same object.

I guess my question is why does pjit._cpp_pjit.cache_miss gets triggered and how can I debug that if at all?

[jakevdp]

_cpp_pjit.cache_miss is called here:

jax/jaxlib/pjit.cc

Line 621 in cf9a352

PyObject_Vectorcall(cache_miss_.ptr(), args, nargs, kwnames));

There are a few contexts in this file where it falls back to a cache miss – e.g. where the input array is not committed and there is more than one global device. It might be that you're hitting one of these conditions.