lax.cond executing both branches?

[chaserileyroberts]

When using lax.cond, is it ever the case that both branches get executed on device, regardless of the value of the cond? Does marking the primitives as having side effects prevent this at all?

[mattjj]

When using lax.cond, is it ever the case that both branches get executed on device, regardless of the value of the cond?

Both Python callables are always executed, in the sense of being applied to abstract values to stage out their contents. But no FLOPs would happen.

I assume you're talking about both sides' FLOPs getting evaluated. One case is that vmap-of-cond basically produces a lax.select (aka a jnp.where), meaning both sides' FLOPs get evaluated. Could that be relevant?

Does marking the primitives as having side effects prevent this at all?

It should, but it doesn't! It's on our todo list to fix, but to do that we probably need to teach JAX's internals about side effects more holistically. #8699 is the first step in that direction, but currently it's just a draft.

[chaserileyroberts]

I assume you're talking about both sides' FLOPs getting evaluated. One case is that vmap-of-cond basically produces a lax.select (aka a jnp.where), meaning both sides' FLOPs get evaluated. Could that be relevant?

This is a good foot gun to know about, but for my specific situation I'm only worried about cond directly as a primitive itself. For context, we're essentially creating our own jaxpr rewrite tool with the rule

rewrite: cond(branch1, branch2) -> cond(rewrite(branch1), rewrite(branch2))

If at FLOP execution time only one of the branches is actually evaluated, then our above rewrite rule would be fine. However, if both get execute at FLOP time, then our rewrite rule becomes undefined and we'd need to warn the user instead.