jax.lax.cond aggressively eliminated at tracing time?

[femtomc]

I read a previous discussion about the runtime cost of jax.lax.cond.

I did a few experiments -- and it seems like this primitive is not aggressively pruned, when the first bool argument is a trace-time constant.

Should I be worried about this at all from an optimization perspective?

[jakevdp]

I'm not sure what experiments you did or what you were measuring, but it is the case that lax.cond will not compute the output for the unused path at runtime, even if the boolean argument is not a constant. Here's a simple example demonstrating this:

import jax.numpy as jnp
import jax

@jax.jit
def f_slow(x):
  M = jnp.outer(x, x)
  evals, evecs = jnp.linalg.eigh(M)
  return evals.sum()

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

@jax.jit
def f_cond(x):
  return jax.lax.cond(x.sum() > 0, f_slow, f_fast, x)

x = jnp.ones(2000)

f_slow(x).block_until_ready()
%timeit f_slow(x).block_until_ready()
# 1.72 s ± 24 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

f_fast(x).block_until_ready()
%timeit f_fast(x).block_until_ready()
# 5.89 µs ± 87.6 ns per loop (mean ± std. dev. of 7 runs, 100000 loops each)

# positive x: should take slow path
f_cond(x).block_until_ready()
%timeit f_cond(x).block_until_ready()
# 1.69 s ± 23.9 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

# negative x: should take fast path
%timeit f_cond(-x).block_until_ready()
# 536 µs ± 10.5 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)

When x is negative, the function runs far faster than it would if the slow branch were being computed.

[femtomc]

@jakevdp -- sorry, I should have specified my question a bit better.

Is the following roughly true:

these branches still appear in the Jaxpr representation, but runtime flow will always choose 1 of the branches (?)

So e.g. I shouldn't expect to see them eliminated from my Jaxpr representation?

[jakevdp]

Yes, I believe both branches of cond will always appear in the jaxpr. Similar to the case of dead code elimination, JAX tends to avoid doing that sort of logic at trace-time, and instead leaves it to the compiler.

[femtomc]

Thank you @jakevdp!