lax.cond, vmap and data dependent while loops may cause non-termination

[lumip]

We recently ran into a situation quite similar to the following (simplified) code:

import jax

def positive_branch(x):
    # this function does not terminate for negative values
    def body(x):
        return x * 0.1

    return jax.lax.while_loop(lambda y: (y <= 0) | (y > 1e-6), body, x)

def negative_branch(x):
    # this function always terminates
    def body(x):
        return x * 0.1

    return jax.lax.while_loop(lambda y: y < -1e-6, body, x)

def branching_fn(x):
    return jax.lax.cond(x >= 0, x, positive_branch, x, negative_branch)

x = jax.random.normal(jax.random.PRNGKey(0), shape=(10,))

branching_fns = jax.vmap(branching_fn)
y = branching_fns(x)

print(y)

So in essence, we have a a while loop that only terminates if the condition that leads the branch it resides in to be executed is true. In classical sequential code this is not a problem, however, in this particular example, it results in non-termination.

What seems to be happening is that despite conditional branching using jax.lax.cond, both branches are executed (we suspect that jax.lax.cond in this case performs some kind of result multiplexing/selection instead of actual conditional executing).

What's a bit puzzling about this is that this does not occur when not using jax.vmap:

• execution on a single data point does not exhibit the problem
• jax.lax.map does not exhibit the problem

A fix in this example is to include the branching condition in the loops conditions to ensure they always terminate - but it might be worthwhile to think about how to fix this in the underlying implementation.. so can you maybe shed some light on what is causing this and whether it could be addressed? At least I think this should be documented somewhere (maybe on the sharp bits page). Also, what exactly is the difference between map and vmap in terms of performance/parallelization?

[jekbradbury]

Most of the time, lax.cond is evaluated using true conditional execution (although the XLA compiler sometimes turns it into lax.select when the body is very small). But our rule for vmap of cond (with a batched condition) unconditionally promotes the cond into select, causing the behavior you're seeing.

Why do we do this? It's related to your question about the difference between map and vmap. While lax.map is implemented with a sequential loop, vmap is a semantically parallel loop typically implemented by running the same code simultaneously on many accelerator cores ("single-program-multiple-data" parallelism). The definition of "the same code" can vary; while some accelerators would be fine with running identical code that includes independent conditionals on each core, others can't do this and require the control flow/instruction stream to be the same between cores, too; vmap (and JAX/XLA's accelerator support in general) targets this more restricted paradigm.

It would be possible for us to maintain our strong SPMD constraint and still run your code correctly if we added a condition everywhere along the lines of "check that at least one instance in the batch should run" before executing the true and false sides; I'm turning your issue into a feature request for that behavior.

[lumip]

Hey @jekbradbury ,
thanks for the insights so far. While your proposed solution could potentially be a general performance improvement, it won't help with my problem because if the data we vmap over contains instances for both branches, we will still end up with negative instances in the positive branch and thus non-termination.

I guess then - with the constraints for vmap as they are - there's no easy way to work around this. So what would be your recommendation to implement the example above in the most performant way? Just use lax.map? You said that's just implemented as a sequential loop.. is there something like a parallelized loop?

[jekbradbury]

You're right, and that's why my attempt in #2961 didn't help. Maybe what we need to do is make the while loop implementation (which already maintains an "execution mask" and terminates once none of its "threads" are active) aware of the higher-level execution mask that the containing cond is introducing.

[jekbradbury]

(For now I would probably do the batching/add the masks manually, inside the condition of the while loop.)

[terhorst]

Has anything happened elsewhere in jax in the past ~15 months that would provide a better solution? I ran into exactly this problem when porting over scipy.special.expi and scipy.special.expn. For positive x, these functions are computed using series approximations and/or continued fractions that terminate according to a convergence criterion. They are extended to negative x via the identity expi(x) = -expn(1,-x). Thus, even using lax.cond to branch on the sign of x, I get the same bug when using vmap whereby the loops used to evaluate the series/CFs do not terminate due to the series itself diverging for negative x.

[soraros]

As a workaround, you could do this:

def body(x):
  return 0.1 * x

def true_fun(arg):
  x, pred = arg
  return lax.while_loop(lambda x: pred & ((x <= 0) | (x > 1e-6)), body, x)

def false_fun(arg):
  x, pred = arg
  return lax.while_loop(lambda x: ~pred & (x < -1e-6), body, x)

def branching(x):
  cond = x > 0
  return lax.cond(cond, true_fun, false_fun, (x, cond))

a = jnp.arange(10.) - 5
vmap(branching)(a)

I guess it's what @jekbradbury meant by adding the mask inside the condition of the while loop.