Conditional Logic with `jax.lax.cond` and `jax.lax.switch`

1. The Clash: Python if vs. JAX jit

You understand that JIT compilation involves a trace-then-compile strategy. The first time a @jit function is called, JAX doesn't run it with the actual numerical values of your inputs. Instead, it traces the function using abstract "tracer" objects that capture the shape and dtype of the arrays, but not their concrete values. This trace generates an intermediate representation (jaxpr) which is then compiled by XLA into an efficient kernel. This kernel can then be reused for any future calls with inputs of the same shape and dtype, avoiding re-compilation overhead.

Herein lies the problem. What happens when the tracer encounters a standard Python if statement whose condition depends on an input value?

if x > 0:

The tracer object for x doesn't have a concrete value, so x > 0 evaluates to another tracer object, a ShapedArray((), jnp.bool_), representing an abstract boolean. The Python interpreter has no way to coerce this abstract boolean into a concrete True or False to decide which branch of the if statement to trace. It's stuck, and throws an error.

Let's see this explained in detail in the JAX documentation.

This video tutorial provides a hands-on demonstration of this exact failure mode.

A Workaround with a Cost: static_argnums

One way to resolve this is to tell JAX to not abstract away the value of the argument used in the condition. You can do this using the static_argnums (or static_argnames) argument to jit. This forces JAX to use the concrete value of that specific argument during tracing.

import jax
import jax.numpy as jnp




# This will fail
@jax.jit
def f_fail(x):
  if x > 0:
    return x * 2
  else:
    return x / 2




# This will work... with a catch
@jax.jit(static_argnums=0)
def f_static(x):
  if x > 0:
    return x * 2
  else:
    return x / 2

print(f_static(10))  # Traces and compiles a version for x > 0
print(f_static(-10)) # Traces and compiles a *new* version for x <= 0

The catch is significant: JAX now generates a new compiled function for every new static value it encounters. If your conditional logic depends on a boolean flag, this is acceptable (two compilations). But if it depends on an integer that can take many values, you'll suffer from constant recompilation, defeating the purpose of jit. This is a workaround, not the idiomatic solution for dynamic runtime conditions.

2. The Idiomatic Solution: jax.lax.cond

The proper JAX-native way to handle conditional logic inside a JIT-compiled function is to use the primitives provided in the jax.lax module. For an if/else block, the primitive is jax.lax.cond.

The signature is: lax.cond(pred, true_fun, false_fun, *operands)

• pred: A boolean scalar JAX array. This is the condition.
• true_fun: A function to be called if pred is True.
• false_fun: A function to be called if pred is False.
• *operands: The arguments that will be passed to either true_fun or false_fun.

Critically, true_fun and false_fun must be pure functions that accept the same number and type of arguments (operands) and return pytrees with the exact same structure. This means the number of return values, their shapes, and their dtypes must match between the two branches. This is a non-negotiable constraint for the compiler to generate a static graph.

Let's rewrite our failing example using lax.cond:

import jax
import jax.numpy as jnp

@jax.jit
def f_cond(x):



  # Define the functions for the true and false branches
  true_branch = lambda op: op * 2
  false_branch = lambda op: op / 2
  



  # Use lax.cond to select the branch based on the predicate
  return jax.lax.cond(x > 0, true_branch, false_branch, x)

print(f_cond(10.0))   # Correctly returns 20.0
print(f_cond(-10.0))  # Correctly returns -5.0
print(f_cond(jnp.array(5.0))) # Also works

This works perfectly inside jit. lax.cond is not a Python if. It's a primitive that gets compiled by XLA into its own conditional instruction. Both branches are traced to ensure their return structures match, but at runtime, only the selected branch is executed. This avoids both the Python tracer error and the recompilation cost of static_argnums.

3. Multiple Branches with jax.lax.switch

What if you have more than one condition, like an if/elif/else chain? For this, JAX provides jax.lax.switch.

The signature is: lax.switch(index, branches, *operands)

• index: An integer scalar JAX array selecting which branch to execute.
• branches: A tuple or list of functions. branches[i] is called if index is i.
• *operands: The arguments passed to the chosen branch function.

Just like with lax.cond, all functions in branches must have identical signatures and return pytree structures.

@jax.jit
def f_switch(idx, x):
  branches = (
    lambda op: op,         # case 0: identity
    lambda op: op ** 2,     # case 1: square
    lambda op: op ** 3      # case 2: cube
  )



  # The default case if idx is out of bounds is the last branch.
  return jax.lax.switch(idx, branches, x)

print(f_switch(0, 10.0)) # 10.0
print(f_switch(1, 10.0)) # 100.0
print(f_switch(2, 10.0)) # 1000.0

jax.lax.switch is the JAX equivalent of a switch statement in C or a match-case in modern Python, designed to be compiled efficiently.

4. Summary: Choosing the Right Construct

The JAX documentation provides a table summarizing which constructs are compatible with jit and grad. It's a useful reference to have.

As you can see, for conditional branching (if), the native Python construct is incompatible with jit when the condition is dynamic. lax.cond is the JIT-compatible and differentiable solution. The table also mentions loop constructs, which we will address in our next lesson.

Conclusion

You have now mastered the fundamental tools for expressing conditional logic in a way that is transparent to JAX's JIT compiler. This is a critical step away from standard Python programming towards a functional, compilable mindset.

Key Takeaways:

• Standard Python if statements fail under jit when the condition is a tracer object because Python cannot resolve the abstract boolean value.
• Using static_argnums is a workaround that forces recompilation for each new static value, which is often inefficient.
• jax.lax.cond is the idiomatic, JIT-compatible primitive for if/else logic. It requires that both branch functions have identical input/output pytree structures.
• jax.lax.switch generalizes this to multiple branches, like a match-case statement, with the same structural constraints on its branch functions.
• These lax primitives are compiled into efficient low-level instructions by XLA, allowing for dynamic, data-dependent decisions within high-performance compiled code.

With conditionals under your belt, the other major category of control flow is iteration. In our next lesson, we will explore how to write efficient, parallelizable loops using jax.lax.scan and jax.lax.fori_loop.