Is a function re-compiled each time we call jax.jit?

[llan-ml]

Hi,

I am a newbie to jax. When I'm reading to the documentation, I'm confused about the caching behavior of jit.

In the caching section, it says that "Avoid calling jax.jit inside loops. Doing that effectively creates a new f at each call, which will get compiled each time instead of reusing the same cached function". However, running the following code only produces one printing side effect:

import jax
def unjitted_loop_body(prev_i):
  print("tracing...")
  return prev_i + 1

def g_inner_jitted_poorly(x, n):
  i = 0
  while i < n:
    # Don't do this!
    i = jax.jit(unjitted_loop_body)(i)
  return x + i

g_inner_jitted_poorly(10, 20)
# output:
WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
tracing...
Out[1]: DeviceArray(30, dtype=int32)

The string "tracing..." is only printed once, and it seems that jit does not trace the function again.

[pharringtonp19]

Interesting remark! As a different example, consider the block of code

import jax 
import jax.numpy as jnp 

from jax.config import config
config.update("jax_log_compiles", 1)

def g(x):
  return x + 1 

def h(x):
  return x + 1 

f = jax.jit(g)
print(f(2))
print(jax.jit(g)(2))
print(jax.jit(h)(2))

With an output of

WARNING:absl:Compiling g (140493413346352) for args (ShapedArray(int32[], weak_type=True),).
WARNING:absl:Compiling h (140493413399040) for args (ShapedArray(int32[], weak_type=True),).
3
3
3

At some point of the "transformation process"(going from the python function call to the compiled function), jax must recognize that it already has a complied version of g. I am not sure at what point it recognizes this but it would be interesting to find out.

[llan-ml]

Calling jax.jit inside loops indeed slows down the running speed. But it seems that jax doesn't re-trace the function and re-uses the caching function.