Cost of evaluating the two branches of jnp.where

[fbartolic]

Consider the following function

import jax.numpy as jnp

@jit 
def f(x):
    cond = g(x)
    return jnp.where(cond, A(x), B(x))

where branch B is at least 1000x more expensive to evaluate than A. My understanding is that in a jit-ed function both branches are evaluated if cond is True but branch B is evaluated with using arrays with abstract values and no floating point operations are executed. If that's the case I'd expect that evaluating f(x) for values where g(x) == True is substantially cheaper than evaluating B(x) but I'm not getting any noticeable speedup. What's going on here?

[soraros]

jnp.where is lowered into select which evaluates both branches. You could try lax.cond.

[YouJiacheng]

jnp.where always concretely evaluate both branches, and the output can be a mixture of two branches if cond is an array.
You can use lax.cond for this case, which only allows scalar condition and is lazy evaluated.
lax.cond do exactly what you want: both branches are traced to jaxpr, while only execute one of them in runtime.

[jakevdp]

That's a good idea - I'm not sure where this would best be documented. Do you have any thoughts?

[YouJiacheng]

Oops, thanks for pointing this out, I haven't known that before! IIRC, this behavior isn't documented, would you add some comments to docs? @jakevdp

[fbartolic]

@jakevdp thanks for pointing that out, I'll have avoid using vmap with this function then. I would also appreciate seeing more documentation on this :).

[YouJiacheng]

Maybe add a note in lax.cond document is enough. and it can be better if also add some note in sharp bits about control flow and vmap.

[soraros]

@jakevdp I'd vote for somewhere in The Sharp Bits - Control Flow.