Performance of lax.cond in JAX vs Python if-else

[ZedongPeng]

I'm experimenting with Structured Control Flow Primitives in JAX. While they provide traceable control flow, I'm noticing that the execution speed seems slower compared to native Python control flow. It's possible I might be using them incorrectly, so I'm curious if anyone has experience with this.

Here's a simple benchmark I ran:

from jax import lax, jit
import time

@jit
def lax_cond(x):
    return lax.cond(x > 500, lambda _: x + 1, lambda _: x - 1, operand=None)

# Non-JIT lax.cond
start_time = time.time()
for i in range(1000):
    res = lax.cond(True, lambda _: i + 1, lambda _: i - 1, operand=None)
end_time = time.time()
print("jax.lax.cond time = ", end_time - start_time)

# JIT-compiled lax.cond
lax_cond(1)  # Trigger JIT compilation
start_time = time.time()
for i in range(10000):
    res = lax_cond(i)
end_time = time.time()
print("jax.lax.cond JIT time = ", end_time - start_time)

# Python if-else control flow
start_time = time.time()
for i in range(10000):
    res = i + 1 if i > 500 else i - 1
end_time = time.time()
print("Python if-else time = ", end_time - start_time)

Output:

jax.lax.cond time =  25.638490200042725
jax.lax.cond JIT time =  0.1420300006866455
Python if-else time =  0.0010020732879638672

Specifically, I'm wondering if lax.cond gets compiled at every iteration in the first test, which could explain the significant slowdown.

Has anyone encountered similar performance issues when using lax.cond or even while_loop? Any insights on optimizing its usage would be much appreciated!

[jakevdp]

See FAQ: benchmarking JAX code for some discussion of this.

With such a small program, what you're measuring here is essentially just JAX's function dispatch overhead, which will always be larger than Python's dispatch overhead. However, in a JIT-compiled JAX program, you incur that overhead once total (when you call the JIT-compiled function), whereas in a Python program, you incur that overhead once per operation. The result is that JAX looks worse in microbenchbarks of a single operation on scalar values, but that doesn't generalize to non-trivial programs.

[ZedongPeng]

Hi @jakevdp,

Thank you for your response! I’d like to ask a more specific question regarding my use case.

In my code, I need to call jax.lax.cond and jax.lax.while_loop multiple times. What’s the best approach to structure the code for optimal performance? For example, would the following pattern be advisable?

@jit
def f():
    jax.lax.cond()
    # or
    jax.lax.while_loop()

for i in range(5):
    f(x)

[jakevdp]

It's hard to say without more details: how do you choose whether to call cond or while_loop in each iteration?

[ZedongPeng]

To simplify the discussion, let's focus on using only while_loop in this case. Between the two implementations below, which one is preferable for performance and clarity? I’m also open to any other suggestions for an optimal design. Any insights for further improvement are greatly appreciated!

@jit
def f():
    def cond():
        pass
    def body():
        pass
    jax.lax.while_loop()

for i in range(5):
    f(x)

vs.

for i in range(5):
    def cond():
        pass
    def body():
        pass
    jax.lax.while_loop()

[jakevdp]

The answer depends on the details. In particular: how many loops are you actually doing? Is it always 5? Is the number of loops statically known, or is the number of loops somehow derived from the result of initial loops? How expensive is a single execution of your while_loop? Is it significantly more than JAX's operation dispatch cost (~100ms) or is it comparable?

[ZedongPeng]

Hi @jakevdp,

Thank you for your suggestion! After running some tests, I found that the following implementation using while_loop significantly outperformed the standard Python while loop, running several times faster.

As I understand it, JAX compiles the cond and body functions when using while_loop, which speeds up the iterations inside the loop. However, if the while_loop needs to be executed multiple times, it seems important to define it outside the outer loop and jit the entire while_loop for optimal performance.

@jit
def f():
    def cond():
        pass
    def body():
        pass
    jax.lax.while_loop()

for i in range(5):
    f(x)

[jakevdp]

Glad you figured it out! But be aware that if you executed the code above as written, then the compiler will have dead-code-elliminated everything within f(), because it has no return value and thus is functionally equivalent to def f(): pass. So if you were expecting the while loop to actually execute at runtime, you'll have to return the value generated by the while_loop call (otherwise your benchmark may not be telling you what you think it's telling you).