Confusion about cache of grad(my_jitted_function)

[langmore]

Given a jitted function f, I can check cache size with f._cache_size(). Now, given df = grad(f), I assume calling df(x) populates cache somewhere. Where and how do I check its size?

@jax.jit
def f(x):
  return x**2
df = jax.grad(f)

df(1.)  # Hopefully this populates some cache.
print(f._cache_size())
==> 0  # Was expecting 1

f(1.)  # Of course this populates cache
print(f._cache_size())
==> 1  # As expected, directly calling f populates cache

[mattjj]

Thanks for the question!

It's a bit complicated, both for fundamental reasons (caching transformed versions of functions given the way JAX's tracing works) and for transient historical-path-dependent reasons.

There are two global jit caches: one for C++ dispatch and the other for dispatching handled by Python. The Python cache is fully general in the sense that it works with all manner of transformations applied to the jitted function, while the C++ cache only works with buffers-in-buffers-out dispatching, and hence not transformation-of-jit cases. (The C++ cache is populated by stealing entries from the Python cache when it can.)

The _cache_size() API (which isn't public AFAIK, is it even documented?) in your example queries the C++ cache only. But because that cache doesn't support the grad-of-jit case, you won't see it increase when you evaluate grad-of-jit calls, as in your example.

To see the cache grow in grad-of-jit calls, you need to dig pretty carefully:

import jax

@jax.jit
def f(x):
  return x**2
df = jax.grad(f)

def hilariously_tricky_cache_size():
  from jax._src import dispatch
  return sum(len(entries) for entries in dispatch._xla_callable.__closure__[1].cell_contents.values())

print(hilariously_tricky_cache_size())  # 0
df(1.)  # Hopefully this populates some cache.
print(hilariously_tricky_cache_size())  # 2

The reason there are two entries is that there's a cached jitted forward pass and a cached jitted backward pass.

If you'd like a public API for that, open a feature request and let's discuss it!

WDYT?

[VolodyaCO]

I am not. I am just testing on my side that jax.clear_backends() works as expected. I need to clear my cache from time to time to avoid out of memory errors. I use the hilariously_tricky_cache_size to check that I don't have anything in cache at all.

[yashk2810]

I think it does work as expected. We have tests for it to make sure it does! The problem is that you are accessing internal APIs here which will break as it did with this release. If you want you can look at pjit._pjit_lower_cached.cache_info but this can break again in the future (if we decide to change the name or something else).

That being said, we will work on a public API.

[VolodyaCO]

Is there any progress here?

[patrick-kidger]

The only solution I've ever found to this is to restart the Python process.

For example I found that the tests for Lineax (run all at once via pytest) ended up OOM'ing. The solution was to run the tests via python -m tests, which triggers this file, which runs pytest separately on each file.

[VolodyaCO]

oh wow. but is there a way right now in the public API to know what the size of the cache is? or is this still something that needs to be done?