How to return 'early' (possibly with an error code) from jitted functions?

[wiep]

I am wondering whether there is a best practice or some syntax sugar to return early from a jit-compatible function. The best I could come up with are nested jax.lax.cond statements. However, especially for longer or complex functions they are not very readable.

Consider the following code as an example. Here, the functions do_work_* should call the functions f, g, and h in a pipe, i.e., forwarding the output from one function to the next. However, if one of the functions called returns an non-zero error code the calling function should return early.

In either case the functions do_work_* will return an object of type Result that holds the error code and the payload. The shape of this object is static given the input. As I said, the jitable function is not very readable and I wonder if there is a better way to achieve the same result.

import jax
import jax.numpy as jnp
from typing import NamedTuple


class Return(NamedTuple):
    payload: jax.Array
    error_code: jax.Array = jnp.int32(0)
    
    @staticmethod
    def error(error_code: int, payload_desc: jax.Array):
        return Return(jnp.empty_like(payload_desc), jnp.asarray(error_code, jnp.int32))

def f(x):
    return 0, x + 1

def g(x):
    return 0, x + 1

def h(x):
    return 0, x + 1



# Function returns early if any of the functions called returns an non-zero
# error code.
def do_work_not_jitable(x: jax.Array) -> Return:
    payload_desc = x

    error_code, fx = f(x)
    if error_code != 0:
        return Return.error(error_code, payload_desc)
    
    error_code, gx = g(fx)
    if error_code != 0:
        return Return.error(error_code, payload_desc)
    
    error_code, hx = h(gx)
    if error_code != 0:
        return Return.error(error_code, payload_desc)
    
    return Return(hx)


# Function returns early if any of the functions called returns an non-zero
# error code but uses jax.lax.cond to avoid python control flow.
def do_work_jitable(x: jax.Array) -> Return:
    payload_desc = x

    def step_h(gx):
        error_code, hx = h(gx)
        return jax.lax.cond(
            error_code != 0,
            lambda _: Return.error(error_code, payload_desc),
            lambda hx: Return(hx),
            hx,
        )

    def step_g(fx):
        error_code, gx = g(fx)
        return jax.lax.cond(
            error_code != 0,
            lambda _: Return.error(error_code, payload_desc),
            step_h,
            gx,
    )
    
    # step_f
    error_code, fx = f(x)
    return jax.lax.cond(
        error_code != 0,
        lambda _: Return.error(error_code, payload_desc),
        step_g,
        fx,
    )


if __name__ == "__main__":
    x = jnp.array([1, 2, 3])
    print(do_work_not_jitable(x))
    print(do_work_jitable(x))

    do_work_jitable = jax.jit(do_work_jitable)
    print(do_work_jitable(x))

[jakevdp]

There's no way to return early from a JIT-compiled function conditioned on a dynamic value. The best available pattern is probably to define parts of your function in blocks, and use lax.cond to string them together; i.e. something like this:

import jax

def f(x):
    return 0, x + 1

def g(x):
    return 0, x + 1

def h(x):
    return 0, x + 1

def identity(x):
  return 0, x

@jax.jit
def do_work(x):
  error_code, x = f(x)
  error_code, x = jax.lax.cond(error_code == 0, g, identity, x)
  error_code, x = jax.lax.cond(error_code == 0, h, identity, x)
  return x

print(do_work(2))
# 5

If you're interested in tracking and returning error codes from your functions, you can certainly do that manually (by adding an extra return value) or you can do so more automatically using the experimental checkify transformation, which you can read about here: https://jax.readthedocs.io/en/latest/_autosummary/jax.experimental.checkify.checkify.html

[wiep]

I was not aware of the module checkify. Thanks for the hint.