jax.vmap without jit compilation traces multiple times?

[EGalahad]

Hi, I would like to ask about the basic concepts of JAX. I’ve noticed that without jit, both vmap and grad seem to re-trace the function each time they are called, and I don’t fully understand why this happens.

From what I understand, JAX traces the batched array to generate the computation graph. However, for an identical input shape passed in subsequent calls, it should be able to reuse the computation graph, rather than re-tracing it every time.

I believe the key difference between using jit and not using it is whether the computation graph is compiled. My assumption was that without jit, the function would pass through an uncompiled version of the graph, while with jit, it would pass through the compiled machine code. However, it seems like the function is traced each time, even without jit.

import jax
import jax.numpy as jnp

# A simple function to be vectorized
def f(x):
    print("f:", x)
    return x ** 2 + 2 * x + 1

# Vectorized version using vmap
print("Test jax.vmap")
f_vmap = jax.vmap(f)

x = jnp.array([1, 2, 3])
print(f_vmap(x))  # Output: [4, 9, 16]

x = jnp.array([2, 2, 3])
print(f_vmap(x))  # Output: [4, 9, 16]


print("\nTest jax.jit(jax.vmap)")
f_vmap_jit = jax.jit(jax.vmap(f))
x = jnp.array([1, 2, 3])
print(f_vmap_jit(x))  # Output: [4, 9, 16]

x = jnp.array([2, 2, 3])
print(f_vmap_jit(x))  # Output: [4, 9, 16]

Output

$ python test_jax_jit.py 
Test jax.vmap
f: Traced<ShapedArray(int32[])>with<BatchTrace> with
  val = Array([1, 2, 3], dtype=int32)
  batch_dim = 0
[ 4  9 16]
f: Traced<ShapedArray(int32[])>with<BatchTrace> with
  val = Array([2, 2, 3], dtype=int32)
  batch_dim = 0
[ 9  9 16]

Test jax.jit(jax.vmap)
f: Traced<ShapedArray(int32[])>with<BatchTrace> with
  val = Traced<ShapedArray(int32[3])>with<DynamicJaxprTrace>
  batch_dim = 0
[ 4  9 16]
[ 9  9 16]

[jakevdp]

Hi - thanks for the question! This is working as expected: outside of JIT, everything is executed eagerly, and there is no cache for tracing functions. This means that when you execute a vmapped function that is not JIT compiled, it will be retraced each time. If this is not the behavior you want, I'd suggest wrapping the vmapped function in JIT to take advantage of the JIT cache. Does that make sense?