jit(cond) compiling prim_fun over and over again?

[cvjjm]

After activating compile logging JAX_LOG_COMPILES=1 I see hundreds of occurrences of what looks like jax tracing and re-compiling the prim_fun of a lax.cond while running a longer jax based computation:

[...]
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.0024445299059152603 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489734464) for with global shapes and types (ShapedArray(int32[]), ShapedArray(float64[96,96])). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.4972473056986928 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.0023525459691882133 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489735344) for with global shapes and types (ShapedArray(int32[]),). Argument mapping: (OpShardingSharding({replicated}),).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.5909602697938681 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.0024695219472050667 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489737104) for with global shapes and types (ShapedArray(int32[]), ShapedArray(float64[96,96])). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.5074702044948936 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.001832389272749424 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489772128) for with global shapes and types (ShapedArray(int32[]),). Argument mapping: (OpShardingSharding({replicated}),).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.5399841964244843 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.0023319674655795097 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489770528) for with global shapes and types (ShapedArray(int32[]), ShapedArray(float64[96,96])). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.38560344371944666 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.0024183234199881554 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489736064) for with global shapes and types (ShapedArray(int32[]),). Argument mapping: (OpShardingSharding({replicated}),).
WARNING  jax._src.dispatch:dispatch.py:382 Finished XLA compilation of jit(cond) in 0.17765467800199986 sec
WARNING  jax._src.dispatch:dispatch.py:382 Finished tracing + transforming jit(cond) in 0.00044081080704927444 sec
WARNING  jax.interpreters.pxla:pxla.py:2822 Compiling prim_fun (22666489770048) for with global shapes and types (ShapedArray(int32[]), ShapedArray(float64[96,96])). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
[...]

There are only very few lax.cond() calls in that code and I was unable to reproduce this behavior in a simple example program.
I first thought that a lax.cond() calling an inline lambda function in a for loop of a jit compiled function could be causing this, but this does not seem to trigger this.

Does anyone know what could be causing this? Since every compilation takes on the order of ~0.5s this is quite a time penalty...

[jakevdp]

I first thought that a lax.cond() calling an inline lambda function in a for loop of a jit compiled function could be causing this, but this does not seem to trigger this.

If you are passing a lambda to a cond within a for loop, this will cause exactly the recompilation behavior you're seeing. Perhaps try defining the function outside the loop rather than inside?

For example, I'd expect this to cause 10 recompilations of each lambda function:

for i in range(10):
  a = lax.cond(x, lambda: a, lambda: a + 1)

Whereas this will produce the same result, but I would expect to see only a single compilation of each lambda function:

same = lambda a: a
augment = lambda a: a + 1
for i in range(10):
  a = lax.cond(x, same, augment, a)

[cvjjm]

I had tried more or less exactly what you suggested in a jit compiled function

@jax.jit
def foo():
    for x in range(10):
        a = lax.cond(x, lambda a: a, lambda a: a + 1, x)

foo()

and then only a single compile happens.
I now see that if I do exactly what you suggested (fixing typos)

for x in range(10):
  a = lax.cond(x, lambda a: a, lambda a: a + 1, x)

or comment out the @jax.jit then I see the recompilation.

Conclusion: Do not use inline functions with jax.cond() in code that is not jit compiled itself!

Thanks!

[cvjjm]

But, defining the prim_fun outside the loop (and even marking it as jit compiled manually) does not solve the issue as this example demonstrates:

import jax
from jax import lax
from jax import numpy as np

@jax.jit
def bar(x):
    return 2*x

@jax.jit
def bar2(x):
    return -x

def foo():
    for x in range(10):
        a = lax.cond(x, bar, bar2, x)

foo()

This re-compiles prim_fun 11(?) times.

[cvjjm]

Why does lax.cond() compile its primary and secondaray function in the first place? Wouldn't it be more efficient if it simply acted like a normal if as long as the contect in which it is used is not jited?

[jakevdp]

I don't see the behavior you're reporting when I run your example. This is the logged output:

WARNING:jax._src.dispatch:Finished tracing + transforming prim_fun for jit in 0.0004899501800537109 sec
WARNING:jax._src.dispatch:Compiling prim_fun (140577066035056 for args (ShapedArray(bool[]),).
WARNING:jax._src.dispatch:Finished XLA compilation of convert_element_type in 0.4547111988067627 sec
WARNING:jax._src.dispatch:Finished tracing + transforming prim_fun for jit in 0.0007920265197753906 sec
WARNING:jax._src.dispatch:Compiling prim_fun (140576874933920 for args (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)).
WARNING:jax._src.dispatch:Finished XLA compilation of cond in 0.06107807159423828 sec

Can you say more about the environment where you're seeing these multiple re-compiles?

[soraros]

@jakevdp It can be reproduced on colab with

jax:    0.4.1
jaxlib: 0.4.1
numpy:  1.21.6
python: 3.8.10 (default, Nov 14 2022, 12:59:47)  [GCC 9.4.0]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1

[jakevdp]

Thanks - looks like the extra compiles are a regression in JAX 0.4.1. I'll file a bug (#14058)

[cvjjm]

Thanks!