Need an input-dependent if statement that still works with JAX transformations? lax.cond selects one of two functions using a predicate, while remaining representable in Jaxpr and compatible with JIT compilation and differentiation. section 3.4.1 section 5.2.1
A traceable ReLU
def relu(x): return jax.lax.cond(x > 0, lambda x: x, lambda x: 0.0, x) relu(12.3) # Array(12.3, dtype=float32, weak_type=True) Unlike a Python if whose condition depends on a traced value, this form produces a cond equation in Jaxpr and can be passed to jax.jit. section 5.2.1
Listing 5.17 Replacing an if statement with a structured control flow primitive
def relu(x):
return jax.lax.cond(x>0, lambda x: x, lambda x: 0.0, x) #1
relu(12.3)
>>> Array(12.3, dtype=float32, weak_type=True) #2
jax.make_jaxpr(relu)(12.3) #3
>>> { lambda ; a:f32[]. let
>>> b:bool[] = gt a 0.0
>>> c:i32[] = convert_element_type[new_dtype=int32 weak_type=False] b
>>> d:f32[] = cond[
>>> branches=(
>>> { lambda ; e:f32[]. let in (0.0,) }
>>> { lambda ; f:f32[]. let in (f,) }
>>> )
>>> linear=(False,)
>>> ] c a
>>> in (d,) }
jax.jit(relu)(12.3) #4
>>> Array(12.3, dtype=float32, weak_type=True)
Inspect the complete replacement of a data-dependent if statement, including its Jaxpr output and jax.jit call. section 5.2.1
Python control structures can work when their decisions are concrete, but traced conditions can raise TracerBoolConversionError. Use lax.cond when the predicate depends on a runtime JAX value; use lax.switch for multiple branches. section 3.4.1 section 5.2.1
