How does jax caching - recompilation work?

[mohamad-amin]

I have a jax.pmap-ed function that takes in some broadcasted and some indexed inputs. let's say:

def a(b, c):
    do something with a and c

and then i have:

parallel_a = jax.pmap(vmap(a, in_axes=(0, None)), in_axes=(0, None))

I call parallel_a with different (as, b), and I'm noticing one thing: as l change dimensions of the variables of the function, there tends to be a huge performance decrease in comparison to when I don't. As an example:

for i in range(10):
    as = np.random.random((100, 100))
    b = np.random.random((100, 10))
    z = parallel_a(as, b)
    print(z)

Following this code, the first call to parallel_a takes 6 seconds, and the next calls take only 0.1 seconds. I imagine this is due to recompiling the function. As another example:

for i in range(1, 11):
    as = np.random.random((100, 10*i))
    b = np.random((10*i, 10))
    z = parallel_a(as, b)
    print(z)

here, all of the 10 calls will take 6 seconds to run.

I'm using 'XLA_PYTHON_CLIENT_PREALLOCATE' = false and 'XLA_PYTHON_CLIENT_ALLOCATOR' = platform

My questions are:

• Is this behaviour expected?
• If yes, is there anyway to change this behaviour and make the second loop run faster?
• If not, what could be the problem with my code that causes it to run like this?

Thanks!

[mohamad-amin]

As a more general question, is there any method to check if on a run, we're gonna use a compiled version or we're gonna compile and then run?

[jakevdp]

Thanks for the question! The full answer is a bit more involved, but in general recompilation is triggered by changes in the static properties of input arrays. So if a JIT-compiled function sees an input with the same shape and dtype as one it has seen before, the cached code is used. If the function sees an input with a shape and/or dtype it has not seen before, this triggers a re-compilation. With this in mind, the recompilations you're seeing are exactly as expected.

Why is this? Mainly, it's due to the way XLA compilation works: programs are compiled for inputs of a particular shape. There has been some amount of work on supporting dynamic shapes in XLA & in JAX, but it hasn't quite made it to maturity.

[cgarciae]

Hey @jakevdp! I think I've posted about this before but it would be awesome if JAX had some tooling that told the user exactly what is causing recompilation. It can get very tricky as debugging changes in shape, dtype, and static components for highly nested structures can be complicated.

[mohamad-amin]

Hey! Thanks. I understand this, but consider this case:
I may have to run my code for the same collection of shapes over and over again, for example:
Func(Shape 1), Func(Shape 2), Func(Shape3), Func(Shape 1), Func(Shape 2), Func(Shape3), Func(Shape 1), Func(Shape 2), Func(Shape3)
In this case, there should be a way to cache the compiled functions so that I can reuse them, otherwise it would be very slow, am I right?
If yes, does jax provide any caching strategy for this?

[jakevdp]

Yes, JAX caches multiple versions of previously compiled functions, so the second time it sees the same input shapes within a single runtime it will use the cached version, even if other arguments have been passed to the function in the meantime. Note that the cache entry uses weak references to the function object, so if you compile a temporary function like jit(lambda x: x) it will be removed from the cache when the original function goes out of scope.

The cache is an LRU cache with a capacity of 4096, so if you have more compiles than this then earlier ones will be dropped.

If you want to inspect the size of the compiled function cache within your runtime, you can do it this way:

from jax._src.api import _cpp_jit_cache
print(_cpp_jit_cache.size())

[mohamad-amin]

Is the cache LRU(4096) for each function or for all the compiled functions? Because I have 20 different shapes that are fed as input one by one and each time it feels like it's recompiling it.