Expected output with jax.config.update("jax_explain_cache_misses", True)?

[cgiovanetti]

I profiled a code I'm working on recently and noticed some cache misses, so I set jax.config.update("jax_explain_cache_misses", True). There are a lot of persistent cache misses during compilation, but there's no output when the compiled code is called (doesn't seem like async is responsible, I'm block_until_readying many segments of the code, though maybe there's something I'm missing).

Is this what I should expect? I'm using the tensorboard profiler and am seeing blocks like $pjit.py:302 cache_miss, $pjit.py:1507 call_impl_cache_miss in calls to the compiled code. I can set up an MWE and open an issue if this is unexpected.

[jakevdp]

When compiled code is called, you're likely hitting the jit cache, so you shouldn't see any cache miss logs at that point. Does that make sense?

[cgiovanetti]

Sorry, I don't fully understand. Should my takeaway be that those blocks labeled cache_miss in the profiler output of the compiled code aren't really cache misses, or aren't the kinds of cache misses that can be logged with this flag?

[jakevdp]

I'm not sure what the cache_miss blocks in the profiler mean, but the jax_explain_cache_misses logs will only trigger during initial compilation, not when the compiled function is called. Does that help?

[cgiovanetti]

Got it, thank you!