Control Flow in JAX: Mastering jax.lax Primitives

Introduction

Welcome to the fifth and final lesson of "JAX Fundamentals: NumPy Power-Up"! You've made tremendous progress throughout this course, and I'm excited to share this concluding lesson with you. So far, we've mastered JAX arrays and their immutable nature, understood the importance of pure functions, explored automatic differentiation with jax.grad, and unlocked dramatic performance improvements with jax.jit compilation. Each of these concepts has been building toward today's sophisticated topic. Today, we're tackling control flow in JAX using the powerful primitives from jax.lax . As you may recall from our previous lesson on JIT compilation, JAX transforms and compiles your functions to achieve remarkable performance gains. However, this transformation process imposes certain constraints on how we can use traditional Python control flow statements like if/else and for loops within JIT-compiled functions. In this lesson, we'll discover why standard Python control flow can be problematic in JAX's compiled context and learn how to use jax.lax.cond, jax.lax.scan, and jax.lax.while_loop to implement conditionals and loops that work seamlessly with JAX's transformations. These primitives aren't just workarounds, as they're often more efficient and expressive than their Python counterparts for numerical computing tasks. By the end of today's lesson, you'll have a complete foundation in JAX fundamentals, ready to tackle more advanced topics in your continued JAX journey.

Why Python Control Flow Can Be Problematic in JIT

As we discussed before, when JAX compiles a function with jax.jit, it doesn't execute your Python code directly. Instead, it performs a process called tracing . During tracing, JAX runs through your function using abstract, symbolic representations of your inputs (called tracers) rather than concrete values. This tracing process captures the computational graph (the sequence of operations your function performs) which JAX then compiles into optimized XLA (Accelerated Linear Algebra) code. The challenge arises with Python's control flow statements. Consider a simple if statement that depends on the value of a JAX array: Python # A Python if statement if x > 0: # do something else: # do something else During tracing, x is not a concrete number but an abstract tracer object representing "some array of a certain shape and type." When JAX encounters x > 0, it cannot evaluate this condition concretely because x doesn't have a specific numerical value during tracing: it's just a placeholder. Python's if statement requires a concrete boolean value (True or False) to decide which branch to execute. However, JAX can only provide an abstract representation of the comparison result (e.g., a JAX boolean array). This fundamental mismatch between Python's eager evaluation model and JAX's deferred, symbolic computation model during tracing is what necessitates special control flow primitives. JAX needs to capture both branches of a conditional or the entire structure of a loop in the compiled function, rather than committing to just one path during tracing. To solve this issue, we can employ the jax.lax module!

Join the 1M+ learners on CodeSignal

Be a part of our community of 1M+ users who develop and demonstrate their skills on CodeSignal

Python

import jax
import jax.numpy as jnp
import jax.lax as lax # We'll use lax later

@jax.jit
def problematic_conditional(x):
    # This Python 'if' statement will cause issues during JIT compilation
    # because 'x > 0' cannot be resolved to a Python boolean at trace time.
    if x > 0: # Condition depends on a JAX array value
        return x * 2
    else:
        return x + 10

# Attempting to call this JIT-compiled function will raise an error
try:
    # We pass a JAX array, so 'x' inside problematic_conditional is a tracer
    result = problematic_conditional(jnp.array(5.0)) 
except Exception as e:
    print(f"Error with Python if/else in JIT: {type(e).__name__}")
    print("JAX can't evaluate concrete conditions during tracing!\n")

text

Error with Python if/else in JIT: TracerBoolConversionError
JAX can't evaluate concrete conditions during tracing!

Python

# Define the functions for the true and false branches
def true_fun(operand): # Function to execute if condition is true
    return operand * 2

def false_fun(operand): # Function to execute if condition is false
    return operand + 10

@jax.jit
def safe_conditional(x):
    # lax.cond(predicate, true_branch_func, false_branch_func, operand(s))
    # 'x > 0' is the predicate.
    # 'true_fun' and 'false_fun' are the functions for each branch.
    # 'x' is the operand passed to the chosen function.
    return lax.cond(x > 0, true_fun, false_fun, x)

# Test with different values
test_values = [5.0, -3.0, 0.0] # Note: 0.0 > 0 is False
for val in test_values:
    result = safe_conditional(jnp.array(val))
    print(f"safe_conditional({val}) = {result}")

text

safe_conditional(5.0) = 10.0
safe_conditional(-3.0) = 7.0
safe_conditional(0.0) = 10.0

Python

# Function to compute cumulative sum using lax.scan
@jax.jit
def cumsum_scan(arr):
    # scan_fn(carry, x) processes one element x from arr
    # 'carry' is the accumulated state from the previous step
    def scan_fn(carry, x):
        new_carry = carry + x # Update the running sum
        # Returns (new_carry_for_next_iteration, output_for_this_iteration)
        return new_carry, new_carry 
    
    # lax.scan(scan_function, initial_carry, array_to_scan_over)
    # The first returned value is the final carry, which we ignore with '_'
    # The second is an array of the 'output_for_this_iteration' from each step
    _, cumulative_sums = lax.scan(scan_fn, 0.0, arr) # Initial carry is 0.0
    return cumulative_sums

arr = jnp.array([1.0, 2.0, 3.0, 4.0, 5.0])
cumsum_result = cumsum_scan(arr) # Can be JIT-compiled
print(f"Array: {arr}")
print(f"Cumulative sum: {cumsum_result}")

text

Array: [1. 2. 3. 4. 5.]
Cumulative sum: [ 1.  3.  6. 10. 15.]

Python

# Fibonacci sequence using scan
def fibonacci_scan(n_terms): # n_terms is the number of Fibonacci numbers to generate
    def fib_step(carry_pair, _): # '_' because we don't use an input array element
        a, b = carry_pair
        # New carry is (b, a+b), output for this step is 'a'
        return (b, a + b), a 
    
    # Initial carry: (0, 1) for F_0, F_1
    # We scan 'None' for the input array 'xs', but specify 'length'
    # This tells scan to run fib_step 'n_terms' times.
    initial_carry = (jnp.array(0), jnp.array(1)) # Ensure JAX types for carry
    _, fib_sequence = lax.scan(fib_step, initial_carry, None, length=n_terms)
    return fib_sequence

# JIT-compile fibonacci_scan. 'n_terms' determines the loop length,
# so it must be a static argument for JIT compilation.
jit_fibonacci_scan = jax.jit(fibonacci_scan, static_argnums=(0,))

fib_result = jit_fibonacci_scan(10) # Get the first 10 Fibonacci numbers
print(f"\nFirst 10 Fibonacci numbers: {fib_result}")

text

First 10 Fibonacci numbers: [ 0  1  1  2  3  5  8 13 21 34]

Python

@jax.jit
def next_power_of_two(x_target):
    # cond_fn(current_val): returns True if loop should continue
    def cond_fn(current_val):
        return current_val < x_target
    
    # body_fn(current_val): computes the next value of current_val
    def body_fn(current_val):
        return current_val * 2
    
    # lax.while_loop(condition_function, body_function, initial_value)
    # The loop starts with '1' and continues as long as current_val < x_target,
    # doubling current_val in each iteration.
    return lax.while_loop(cond_fn, body_fn, jnp.array(1)) # Start with 1

test_nums = [10, 63, 100]
for num in test_nums:
    result = next_power_of_two(jnp.array(num)) # Pass JAX arrays
    print(f"Next power of 2 after {num}: {result}")

text

Next power of 2 after 10: 16
Next power of 2 after 63: 64
Next power of 2 after 100: 128

Introduction

Why Python Control Flow Can Be Problematic in JIT

Join the 1M+ learners on CodeSignal

Understanding the jax.lax Module

The Problem with Standard Control Flow in Action

Conditional Operations with jax.lax.cond

Loops with Carried State: jax.lax.scan

Fibonacci Sequence with jax.lax.scan

Dynamic Loops with jax.lax.while_loop

Conclusion and Next Steps