jax return 0 if condition, continue if not in a jitted function

Question

I want to replicate this behaviour in a jitted function (the function is an example):

def function(x,y):
   if y==0:
      return x
   return x+1

Using jax.lax.cond it can be obtained with:

@jax.jit
def function(x,y):
   return jax.lax.cond(y==0, lambda x: x, lambda x: x+1, x)

This is simple as long as whatever needs to be done is simple when y!=0 (in this case, just add 1 to x). However, if that's complex, or there are more conditions of this sort, coding gets more convoluted.

Is there a way to get the behavior "if y==0, return x, if not, just keep running the function. jax.lax.cond requires a new function for every condition that is applied.

For example, this starts to become convoluted.

def function(x,y):
    if y==0:
       return x
    if y>0:
       return x-y
    if y<0:
       return x+y

This starts to be messy:

@jax.jit
def function(x,y):
    jax.lax.cond(y==0, 
             lambda x,y: x, 
             lambda x,y: jax.lax.cond(x>0, lambda x,y:x-y, lambda x,y: x+y, x,y),
             x,y)

Is there a better way?

Answer 1

In short, no, there's no way to return early from a Python function conditioned on traced values. The pattern I typically see to avoid messy nesting is to encapsulate logic in helper functions and call them via lax.cond.

Alternatively, if you are branching based on multiple conditions, you may be able to better express the logic in terms of lax.switch; for example:

@jax.jit
def function(x, y):
  branches = [lambda: x, lambda: x-y, lambda: x+y]
  conditions = jnp.array([y == 0, x > 0, True])
  return lax.switch(jnp.argmax(conditions), branches)