Slow jit when having inpure function

[tetterl]

I'm trying to compute a Jacobian-Matrix-Product (J @ M) where the matrix M always stays the same but the Jacobian J varies (in the following example even the Jacobian stays fixed). Thus I thought I can capture the matrix M in the function (aka let JAX copy the matrix into the jitted function). This matrix will never change and therefore I shouldn't get any strange side effects. Nevertheless, this results in a really slow jit time. When providing the the matrix (M) to the function and thus making it a pure function jitting runs as fast as expected. See the following code snippet:

import time

import jax.numpy as jnp
from jax import jit, jvp, random, vmap

key = random.PRNGKey(0)
J = random.uniform(key, (10000, 100))


def f(x):
    return J @ x


x = jnp.ones(J.shape[1])
M = random.uniform(key, (100, 50))

@jit
def vmap_jmp_inpure(x):
    _jvp = lambda s: jvp(f, (x,), (s,))[1]
    return vmap(_jvp, in_axes=1, out_axes=1)(M)

start = time.time()
AM2 = vmap_jmp_inpure(x)
print('inpure jit time: ' + str(time.time() - start))

@jit
def vmap_jmp_pure(x, M):
    _jvp = lambda s: jvp(f, (x,), (s,))[1]
    return vmap(_jvp, in_axes=1, out_axes=1)(M)

start = time.time()
AM1 = vmap_jmp_pure(x, M)
print('pure jit time: ' + str(time.time() - start))

JM = J @ M
assert jnp.allclose(JM, AM1)
assert jnp.allclose(JM, AM2)

Resulting in:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
inpure jit time: 1.500485897064209
pure jit time: 0.022991180419921875

Can somebody shed some light on the reasons why this is so much slower in the case where I (mistakenly) used an impure function?

[tetterl]

Same happens if we jit the partial function:

vmap_jmp_pure_partial = jit(partial(vmap_jmp_pure, M=M))
start = time.time()
AM3 = vmap_jmp_pure_partial(x)
print('pure jit partial time: ' + str(time.time() - start))

resulting in:

pure jit partial time: 1.4687762260437012

Is this a performance bug or at least something that should trigger a warning? (@jakevdp)