When a jit function calling another jit function multiple times, does the latter need to be compiled each time?

[green-slime]

Suppose we've defined a jit function A:

@jax.jit
def A():
    ...

As we know, when a non-jit function B calls A multiple times:

def B():
    A()
    A()

starting from the second invocation of A(), A() will benefit from the acceleration of jit.

But if a jit function B calls A multiple times:

@jax.jit
def B():
   A()
   A()

when first time calling B(), I found that each A() would be compiled, which take a long time.
So, I wonder if is there any way when first calling jit(B), to benefit from the acceleration of jit(A) after the second invocation of A()?

Thank you so much for any help!

[jakevdp]

Thanks for the question!

When I try the example you give, A appears to only compile/trace once:

import jax

@jax.jit
def A():
  print('compiling a')
  return

@jax.jit
def B():
  A()
  A()

B()
# compiling a

print(A._cache_size())
# 1

It is possible that your simple example leaves out relevant details: for examle if the arguments to A have different shape or dtype on the second call, it will result in a cache miss and an additional compilation.

Can you try to put together a minimal reproducible example of the behavior you're seeing, including how you're determining that A is being compiled twice? Thanks!

[green-slime]

Thanks for the question!

When I try the example you give, A appears to only compile/trace once:

import jax

@jax.jit
def A():
  print('compiling a')
  return

@jax.jit
def B():
  A()
  A()

B()
# compiling a

print(A._cache_size())
# 1

It is possible that your simple example leaves out relevant details: for examle if the arguments to A have different shape or dtype on the second call, it will result in a cache miss and an additional compilation.

Can you try to put together a minimal reproducible example of the behavior you're seeing, including how you're determining that A is being compiled twice? Thanks!

Thank you so much for your reply!
I've tried this code on my device:

@jax.jit
def A():
  print('compiling a')
  return

@jax.jit
def B():
  A()
  A()

B()
print(A._cache_size())

However the result of the code is

# compiling a
# compiling a
# 1

Therefore, I infer that A() has been compiled twice.

I guess this might be related to the version of JAX that I am using? I am using jax==0.3.14 and jaxlib==0.3.14+cuda11.cudnn805.

[jakevdp]

Yes, it's possible that the compilation cache worked differently when jax v0.3.14 was released (June 2022) so its not surprising that you're seeing different results than you would with a more recent JAX version.

I'd suggest updating JAX to a more recent release; we're not going to be able to help you debug the details of the compilation cache behavior from such an old JAX version, as I suspect you're hitting an issue that has already been debugged and fixed in a more recent release.

[green-slime]

Yes, updating jax to v0.4.7 solved this problem. Thank you once again for your help!