Lazy evaluation for jnp.select?

[rkruegs123]

Hi, I have a question on the usage of jnp.select.

Consider I do the following:

def f(x):
    return f(x)

conds = [True]
choices = [0.0]
jnp.select(conds, choices, default=f(0))

where f(x) will always cause an infinite loop. Of course this isn't my use case, but it's representative of the problem that I'm having.

Now, since the condition is true, I would expect the infinite loop to never be entered. However, it seems to be evaluated regardless. Is there any way to have lazy evaluation of the default condition? In other words, I would like f(x) to only be evaluated when all conditions are false (of course I am using some f that won't enter an infinite loop under these conditions).

Otherwise, can anyone recommend a good way to implement an if ... elif ... else control flow that has this property but isn't just a bunch of nested lax.conds? A bunch of nested lax.conds is just a bit messy...

Thank you!

[jakevdp]

The general answer to your question is that conditionals based on lax.cond will evaluate lazily, while conditionals based on lax.switch will evaluate every branch regardless of the condition. Since jnp.where lowers to lax.switch, it will always evaluate both branches.

Now there is another issue in your example, which is that it has an infinite recursion at trace-time, so the computation never has a chance to be lowered to XLA.

With that in mind, here's an example of a conditional that will only be evaluated for a single branch at runtime:

from jax import lax
import jax.numpy as jnp

def truefunc(x):
  return x

def falsefunc(x):
  return 100 * x

x = jnp.arange(5)

lax.cond(False, truefunc, falsefunc, x)
# DeviceArray([  0, 100, 200, 300, 400], dtype=int32)

However, keep in mind that both functions will be traced, so you won't be able to check this by inserting a Python-side infinite loop.

Another detail to keep in mind is that if you vmap over lax.cond, it will lower to lax.select, meaning the batched computation will evaluate both branches.

[rkruegs123]

Thank you for your quick and thorough reply. I understand it mostly, but...

Does that mean that the example that you gave with

def falsefunc(x):
    return falsefunc(x)

can't be JIT-ed because it has an infinitne recursion at trace time?

I ask all of this because I am trying to implement a dynamic programming algorithm, and there is a branch point that looks something like the following:

if already_computed:
    return
else:
    compute and store

If both branches are always evaluated during trace time, you get an infinite loop in the path where the "compute and store" branch is always chosen

[jakevdp]

Right, both branches will always be abstractly evaluated at trace time, so you cannot use JAX to evaluate Python control flow in a lazy manner.

[rkruegs123]

Also, incase it isn't clear, I am implementing my DP algorithm with backwards recursion. I would expect this to only be an issue with backwards recursion.

However, if the above property is true, it seems as though a DP algorithm could never be implemented in JAX via backwards recursions, and only with forwards recursions.

[jakevdp]

For more background on JAX's tracing/computation model, you might start with https://jax.readthedocs.io/en/latest/notebooks/thinking_in_jax.html

[YouJiacheng]

However, if the above property is true, it seems as though a DP algorithm could never be implemented in JAX via backwards recursions, and only with forwards recursions.

@sharadmv
I think this is another use case for no-inline/sub-jaxpr/subgraph, use call jaxpr/XLA call (under user's command) instead of tracing through everything will make it possible to write such recursion algorithm.