how to use control flow under jit?

[yiminghwang]

Hi, there,

I have a problem when I want to use jax.jit to speedup my codes, the simplified example is shown below,

    class A:

        @partial(jit, static_argnums=(0, 1))
        def funcA(self, x):
            #some processing codes here, suppose it is y=x*0.5
            y = 0.5*x
            p =  jnp.float32(y) # here the p is a DeviceArrary(()), generated by the above process
            t = self.accept_or_not(p)
            return t
        
        def accept_or_not(self, p):
            uni = jnp.asarray(np.random.uniform(low=0, high=1))
            if jnp.greater(p,uni) == True: # if uni < p:  #---------------->Here it raises an error
                return True
            else:
                return False

    a = A()
    result = a.accept_or_not(jnp.float32(0.6))

In class A, I want to speed up the funcA by jax.jit, it raises an error about "Abstract tracer value encountered where concrete value is expected", how can I fix it?

[jakevdp]

You can replace your if statement with a where statement:

def accept_or_not(self, p):
  uni = jnp.asarray(np.random.uniform(low=0, high=1))
  return jnp.where(uni < p, True, False)

or, even more simply for this particular case (since you're just returning True or False you could do this:

def accept_or_not(self, p):
  uni = jnp.asarray(np.random.uniform(low=0, high=1))
  return uni < p

[yiminghwang]

Thanks for your quick response!
to go a step further, if I have further processing, should I use lax.cond or something else? i.e.

def accept_or_not(self, p):
  uni = jnp.asarray(np.random.uniform(low=0, high=1))
  return jnp.where(uni < p, True, False)

@partial(jit, static_argnums=(0, 1))
def funcA(self, x):
    #some processing codes here, suppose it is y=x*0.5
    y = 0.5*x
    p =  jnp.float32(y) # here the p is a DeviceArrary(()), generated by the above process
    if self.accept_or_not(p):
        #some processing codes here, and have a result.
    return result

for example

operand = jnp.array([0.])
result = lax.cond(self.accept_or_not(p), lambda x: x+1, lambda x: x-1, operand)

[jakevdp]

Yes, in general the pattern is to use lax.cond or jnp.where in place of if statements that depend on traced variables. See https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#python-control-flow-jit for some more discussion of this.

[yiminghwang]

Thanks a lot!