Control Flow in JAX: Mastering jax.lax Primitives | CodeSignal Learn

Control Flow in JAX: Mastering jax.lax Primitives | CodeSignal Learn

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!

Sign up

Join the 1M+ learners on CodeSignal

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

The Problem with Standard Control Flow in Action

Let's see this control flow issue firsthand. We'll define a function that uses a standard Python if/else statement and attempt to JIT-compile it. 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") When you run this code, you'll observe the following output: Error with Python if/else in JIT: TracerBoolConversionError JAX can't evaluate concrete conditions during tracing! The TracerBoolConversionError clearly indicates the problem: JAX is trying to convert a tracer (the symbolic representation of x > 0) into a Python boolean, which isn't possible during the tracing phase of JIT compilation. JAX cannot determine which branch (x * 2 or x + 10) to embed into the compiled code because the condition's outcome is unknown at compile time. This is where JAX's structured control flow primitives come to the rescue.

Conditional Operations with jax.lax.cond

For conditional logic that depends on JAX array values within JIT-compiled functions, JAX provides jax.lax.cond. This primitive allows JAX to trace and compile both potential execution paths. Let's rewrite our previous example using jax.lax.cond: # 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}") Running this corrected code yields: safe_conditional(5.0) = 10.0 safe_conditional(-3.0) = 7.0 safe_conditional(0.0) = 10.0 Here's a breakdown of jax.lax.cond(pred, true_fun, false_fun, *operands): pred: The boolean condition (a JAX scalar boolean array). true_fun: A callable (function) that is executed if pred is True. It takes *operands as arguments. false_fun: A callable (function) that is executed if pred is False. It also takes *operands as arguments. *operands: The arguments passed to either true_fun or false_fun. Crucially, both true_fun and false_fun must have the same JAX type signature; that is, their output arrays must match in terms of same shapes and dtypes. JAX traces both functions, and the compiled code will efficiently select which branch to execute at runtime based on the concrete value of the predicate.

Loops with Carried State: jax.lax.scan

For loops where each iteration depends on a state carried over from the previous one (like a for loop with an accumulator), JAX offers jax.lax.scan. This is highly efficient for operations like cumulative sums, running averages, or implementing recurrent neural networks. Let's implement a cumulative sum using jax.lax.scan: # 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}") This will output: Array: [1. 2. 3. 4. 5.] Cumulative sum: [ 1. 3. 6. 10. 15.] The lax.scan function takes: f: The scan function, (carry, x) -> (new_carry, y). init: The initial carry value. xs: The array to iterate over (each element x is passed to f). It returns (final_carry, ys), where ys is the stack of y values from each iteration.

Fibonacci Sequence with jax.lax.scan

We can also use lax.scan for more complex sequences, like Fibonacci numbers. Here, the carry will be a tuple (a, b) representing the last two Fibonacci numbers. # 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}") The output for the Fibonacci sequence is: First 10 Fibonacci numbers: [ 0 1 1 2 3 5 8 13 21 34] Notice how n_terms was marked static using static_argnums=(0,). This is because the number of iterations in lax.scan (when xs is None and length is provided) defines the unrolled structure of the computation graph, which needs to be known at compile time.

Dynamic Loops with jax.lax.while_loop

When the number of loop iterations isn't known beforehand and depends on a condition evaluated within the loop (akin to a Python while loop), jax.lax.while_loop is the appropriate tool. Let's use it to find the smallest power of 2 that is greater than or equal to a given number x: @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}") This code will produce: Next power of 2 after 10: 16 Next power of 2 after 63: 64 Next power of 2 after 100: 128 The jax.lax.while_loop takes: cond_fun: A function (val) -> bool_scalar that returns True to continue looping. body_fun: A function (val) -> new_val that executes the loop body. init_val: The initial state val of the loop. Both cond_fun and body_fun operate on the loop state val. The loop state can be any Pytree, but its structure must remain consistent across iterations. jax.lax.while_loop allows JAX to compile the looping logic even when the number of iterations is dynamic.

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")

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}")

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}")

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}")

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}")