jit(grad(f)) isn't cached?

[josephrocca]

Is it expected behavior that repeatedly calling jit(grad(f))(a) will trigger a recompilation each time? I'd have thought that it would work like jit(f) in that repeated calls will use the cache. Below is some example code (and here's a colab link for it - change runtime to GPU). As you can see, jit_grad_f(a) is significantly faster than jit(grad(f))(a).

!pip install --upgrade pip
!pip install --upgrade "jax[cuda111]" -f https://storage.googleapis.com/jax-releases/jax_releases.html

import jax
from jax import grad, jit
import jax.numpy as np

f = lambda x: np.mean(np.dot(x, x.T))
jit_grad_f = jit(grad(f))

# cache/warm-up
a = np.ones([2000,2000])
jit(grad(f))(a)
jit_grad_f(a)

%timeit jit(grad(f))(a).block_until_ready()
# 28.2 ms ± 351 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit jit_grad_f(a).block_until_ready()
# 5.51 ms ± 226 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[cgarciae]

You need to call .block_until_ready() to benchmark properly:

https://jax.readthedocs.io/en/latest/async_dispatch.html

[josephrocca]

Ah, true, thanks! I had that in the real code but forgot to add it when creating this minimal example. I've updated the original post/colab and re-ran the %timeits, but the numbers are still roughly the same:

%timeit jit(grad(f))(a).block_until_ready()
# 32.6 ms ± 678 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

%timeit jit_grad_f(a).block_until_ready()
# 5.67 ms ± 27.9 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)

[cgarciae]

To answer the real question: since jit(grad(f)) creates a new function object every time then it has to compile (everytime) because doesn't have anything in its cache, jit_grad_f does have things in cache because it persists, therefore its faster.

[josephrocca]

@cgarciae Ahh, so you mean that grad(f) isn't cached, and thus returns a new function each time, and so jit must recompile it every time, since it caches compilations based on the function reference rather than the function source string. Makes sense - thanks for your help! Pre-computing grad(f) results in a JIT cache hit as expected:

grad_f = grad(f)

%timeit jit(grad_f)(a).block_until_ready()
# 5.34 ms ± 199 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)