Creating computational graph using an array Bools

[ToshiyukiBandai]

Hi all,

I would like to implement a PDE solver with complicated boundary conditions. To do that, I want to use an array of bools to define computational graph rather than using lax.switch because lax.switch needs to trace all conditions. The example below works for grad but not for jit. Is there any way around?

import jax
import jax.numpy as jnp
from jax import jit, grad

def f(x, flags):
    temp = 0.0
    for i in range(len(flags)):
        if flags[i]:
            print('True!')
            temp = temp + x
        else:
            print('False!')
            pass
    return temp
    

x = jnp.array(1.0)
flags = jnp.array([True, True])

f(x, flags) # works

grad(f)(x, flags) # works

jit(f, static_argnums = (1))(x, flags) # not working!

[jakevdp]

There is no avoiding the tracing of all branches if you want to JIT-compile your branching logic.

If you use lax.switch and don't use vmap, then only one branch will be executed at runtime. But if tracing only a single branch is your concern, then you have to do the branching outside JIT.

[ToshiyukiBandai]

Thank you so much for your quick response! I makes sense. For those who are interested, here is an example showing what happens:

import jax
import jax.numpy as jnp
from jax import jit, grad

def f(x, flags):
    temp = 0.0
    for i in range(len(flags)):
        temp = jax.lax.switch(flags[i], [branch_1, branch_2], temp, x)
    return temp

def branch_1(temp, x):
    print("compiled branch_1!")
    jax.debug.print("Computing branch_1")
    return temp + x
def branch_2(temp, x):
    jax.debug.print("Computing branch_2")
    print("compiled branch_2!")
    return temp

x = jnp.array(1.0)
flags = jnp.array([0, 0])

jitted_f = jit(f)

jitted_f(x, flags)

jitted_f(x, flags)