jnp.where leads to infinite loop?

[dui1234]

I have a tree-like data structure and want to do some operations over the tree. However, when I use jnp.where over some conditions, e.g. condition1 as return jnp.where(condition1, tree[node], body1(mid)) it returned me an error which causes by an infinite loop.

import jax.numpy as jnp
from jax import jit, vmap, lax

def construct_tree(size, init_value, alpha: float = 0.6):
    tree_capacity = 1
    while tree_capacity < size: tree_capacity *= 2
    return [jnp.array([init_value for _ in range(2 * tree_capacity)]),
            tree_capacity, 
            size, 
            alpha, 
            1.0, 
            0]
@jit
def set_tree_value_add(tree):
    idx = tree[-1] + tree[1]
    tree[0] = tree[0].at[idx].set(tree[-2] ** tree[3])
    idx //= 2
    cond_fun = lambda x: x[0]>= 1
    def body_fun(idx_t):
        tree = idx_t[1].at[idx_t[0]].set(idx_t[1][2 * idx_t[0]] + idx_t[1][2 * idx_t[0] + 1])
        idx = idx_t[0]//2
        return [idx,tree]
    dmy = lax.while_loop(cond_fun, body_fun, [idx,tree[0]])
    tree[0], tree[-1] = dmy[1], (tree[-1] + 1) % tree[2]
    return tree

def operate_sum1(tree, start, end):
    end = lax.cond(end <= 0, lambda x: x+tree[1], lambda x: x-1, end)
    return operate_helper_add1(tree[0],start, end, 1, 0, tree[1] - 1)

def operate_helper_add1(tree, start, end, node, node_start, node_end):
    condition1 = start == node_start and end == node_end
    
    print(f"here, {start,end,node,node_start,node_end}")
    mid = (node_start + node_end) // 2
    
    def body1(mid2):
        condition2 = end <= mid2
        print(f"here_m, {mid2} {condition2}")
        def body2(mid3):
            condition3 = mid3 + 1 <= start
            print(f"here_m2, {mid3} {condition3}")
            return jnp.where(condition3,
                             134,
                             operate_helper_add1(tree, start, mid3, 2 * node, node_start, mid3)) + /
                             operate_helper_add1(tree,mid3+1, end, 2 * node+1, mid3 +1 , node_end)) 
        return jnp.where(condition2,
                         mid2+10,
                         body2(mid2))
    print(f"HEEEEERRRR {condition1}")
    
    return jnp.where(condition1, tree[node], body1(mid))

trees = construct_tree(10, 0.0, 0.2)

for i in range(18): trees = set_tree_value_add(trees)

operate_sum1(trees ,0, 10 - 1)

A part of the results from running the above code shows that even though the condition1 is true, the body1 is still considered:

Just wondering what is happening here.

[JiaYaobo]

jnp.where evaluate all branches first while np.where doesn't, here's minimal example

import jax.numpy as jnp
import numpy as np
from jax import jit

@jit
def f(x):
    print('evaluate here!')
    return x+1

jnp.where(True, 0, f(1))
print('----')
np.where(True, 0, f(1))

Output is

evaluate here
----

For this case, you can use lax.switch I think :)

[jakevdp]

lax.switch will have the same issues as jnp.where; see my other answer.

[jakevdp]

Thanks for the question!

The problem here is that you cannot use JAX runtime conditionals like jnp.where, lax.switch, lax.cond etc. in a recursive fashion. The reason is that JAX's tracing mechanism must abstractly evaluate the output, and there's no way for this abstract evaluation to terminate if the termination condition is dynamic.

So when operate_helper_add1 recursively calls itself, it necessarily leads to an infinite loop. I would suggest finding a way to express your computation that does not involve recursion, perhaps using lax.while_loop instead.

If you want to use recursion in JAX, you can only do so using static (i.e. trace-time) conditionals. For example, something like this would work:

@partial(jax.jit, static_argnames=['i'])
def f(x, i):
  if i <= 0:
    return x
  return f(i * x, i - 1)

f(1.0, 10)

but something like this, though it looks equivalent, leads to an infinite loop during tracing:

@jit
def f(x, i):
  return jnp.where(i <= 0, x, f(i * x, i - 1))

[dui1234]

Very much obliged for the answer, it seems work now with one return for condition3:

def body1(dmy):
    dmy[-1] = (dmy[3]+dmy[4]) // 2
    condition2 = lambda x: x[1] <= x[-1]
    def body2_t(dmy):
        dmy[2] = 2*dmy[2]
        dmy[4] = dmy[-1]
        return dmy
    def body2_f(dmy):
        condition3 = lambda x: x[-2] + 1 <= x[0]
        def body3_t(dmy):
            dmy[2] = 2*dmy[2] +1
            dmy[3] = dmy[-2] + 1
            return dmy
        def body3_f(dmy):
            dmy[1] = dmy[-1]
            dmy[2] = 2*dmy[2] 
            dmy[4] = dmy[-1]
            return dmy
        return lax.cond(condition3(dmy),body3_t,body3_f,dmy) ####one return
    return lax.cond(condition2(dmy),body2_t,body2_f,dmy)


dmy_params = [0, 8, 1, 0, trees[1]-1, 0]
condition1 = lambda x: (x[0] == x[3]) & (x[1] != x[4])

lax.while_loop(condition1,body1,dmy_params)

However, my aims is to get a summation if a part where the segments are overlapped, i.e.

def body2(mid3):
            condition3 = mid3 + 1 <= start
            return jnp.where(condition3,
                             134,
                             operate_helper_add1(tree, start, mid3, 2 * node, node_start, mid3)) + /
                             operate_helper_add1(tree,mid3+1, end, 2 * node+1, mid3 +1 , node_end))  ####summation return 

How can I incooperate this point into the above code since the lax.while_loop only works for one return operator......

[jakevdp]

That's a non-static condition, so you cannot use recursion. With some work I bet you could express it as a while loop.