Peeking Inside: Profiling and Debugging JAX

Introduction

Welcome to the fifth and final lesson of our course, "Advanced JAX: Transformations for Speed & Scale"! It's wonderful to have you here as we capstone your advanced JAX journey. Over the past lessons, you've delved into JAX's functional randomness, mastered automatic vectorization with jax.vmap, explored parallel computing using jax.shard_map, and adeptly handled complex data with PyTrees. These are powerful tools, and you've made excellent progress in understanding them. Now, we turn to an essential skill for any proficient JAX developer: profiling and debugging . As your JAX applications grow in complexity, understanding their performance characteristics, identifying bottlenecks, and resolving issues within JIT-compiled functions become paramount. JAX's unique compilation model and asynchronous execution necessitate specialized techniques distinct from standard Python practices. In this lesson, you'll learn the key methods for analyzing performance and debugging JAX code. We'll explore why regular Python print() statements can be misleading in JIT-compiled functions and how jax.debug.print offers a reliable alternative. You'll also master accurate timing of JAX operations using jax.block_until_ready(), learn to create reusable timing decorators for cleaner code, and get a brief introduction to JAX's built-in profiler. These skills will empower you to build efficient, robust, and scalable JAX applications.

Understanding JAX's Asynchronous Execution

Before diving into specific profiling and debugging tools, it's crucial to grasp a core concept of JAX's execution model: asynchronous dispatch . This behavior is fundamental to why JAX requires particular approaches for timing and inspecting code. When you execute a JAX operation, like jnp.dot(x, y), JAX doesn't typically perform the computation and wait for it to finish before moving on. Instead, it uses asynchronous dispatch : the operation is added to a queue for execution on an accelerator (like a GPU or TPU), and control immediately returns to your Python program. The JAX array you get back is essentially a "future" or a "promise" — the actual computation might still be in progress on the device. This asynchronous nature is a key performance feature. It allows your Python code to continue running and queue up more work while the accelerator processes computations in parallel. However, this introduces challenges for: Accurate Timing : If you use standard Python timing tools like time.perf_counter() around a JAX operation, you might only measure the time it took to enqueue the operation, not the actual time spent on computation by the device. Debugging with Print Statements : Inside JIT-compiled functions, regular Python print() statements execute during the tracing phase . This is when JAX analyzes your function to build a computation graph, not during the actual execution with concrete data. Recognizing this asynchronous model is vital for effective JAX development. The techniques we'll cover, such as jax.block_until_ready() for timing and jax.debug.print for debugging, are designed to work correctly within this asynchronous framework.

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 time

# Create some test data for timing
key = jax.random.key(303)
size = 200
mat_x = jax.random.normal(key, (size, size))
key, _ = jax.random.split(key)
mat_y = jax.random.normal(key, (size, size))

import jax
import jax.numpy as jnp
import time

# Create some test data for timing
key = jax.random.key(303)
size = 200
mat_x = jax.random.normal(key, (size, size))
key, _ = jax.random.split(key)
mat_y = jax.random.normal(key, (size, size))

Python

@jax.jit
def matrix_computation(x, y):
    for _ in range(5):  # Simulate some computational work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

@jax.jit
def matrix_computation(x, y):
    for _ in range(5):  # Simulate some computational work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

Python

# Correct timing approach
start_time = time.perf_counter()
result = matrix_computation(mat_x, mat_y)
result.block_until_ready()  # Crucial: Wait for computation to complete on device
end_time = time.perf_counter()
accurate_time = end_time - start_time
print(f"Accurate execution time (first run, includes compilation): {accurate_time:.6f}s")

# Correct timing approach
start_time = time.perf_counter()
result = matrix_computation(mat_x, mat_y)
result.block_until_ready()  # Crucial: Wait for computation to complete on device
end_time = time.perf_counter()
accurate_time = end_time - start_time
print(f"Accurate execution time (first run, includes compilation): {accurate_time:.6f}s")

Python

import functools

def timer(func):
    """A decorator for timing function execution, with JAX support."""
    @functools.wraps(func) # Preserves original function metadata
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        # For JAX arrays, block until computation is complete for accurate timing
        if hasattr(result, 'block_until_ready'):
            result.block_until_ready()
        end_time = time.perf_counter()
        run_time = end_time - start_time
        print(f"Execution took {run_time:.6f}s")
        return result, run_time # Return both result and runtime
    return wrapper

import functools

def timer(func):
    """A decorator for timing function execution, with JAX support."""
    @functools.wraps(func) # Preserves original function metadata
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        # For JAX arrays, block until computation is complete for accurate timing
        if hasattr(result, 'block_until_ready'):
            result.block_until_ready()
        end_time = time.perf_counter()
        run_time = end_time - start_time
        print(f"Execution took {run_time:.6f}s")
        return result, run_time # Return both result and runtime
    return wrapper

Python

@timer # Apply our custom timer
@jax.jit
def timed_function(x, y):
    for _ in range(5):  # Simulate some work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

# First call (includes JIT compilation time)
_, compile_plus_exec_time = timed_function(mat_x, mat_y)
print(f"Time for first run (compilation + execution): {compile_plus_exec_time:.6f}s")

# Subsequent calls (execution only, JIT cache is used)
num_runs = 5
total_exec_time = 0
for _ in range(num_runs):
    _, run_time = timed_function(mat_x, mat_y)
    total_exec_time += run_time
print(f"Average time for {num_runs} subsequent runs: {total_exec_time / num_runs:.6f}s")

@timer # Apply our custom timer
@jax.jit
def timed_function(x, y):
    for _ in range(5):  # Simulate some work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

# First call (includes JIT compilation time)
_, compile_plus_exec_time = timed_function(mat_x, mat_y)
print(f"Time for first run (compilation + execution): {compile_plus_exec_time:.6f}s")

# Subsequent calls (execution only, JIT cache is used)
num_runs = 5
total_exec_time = 0
for _ in range(num_runs):
    _, run_time = timed_function(mat_x, mat_y)
    total_exec_time += run_time
print(f"Average time for {num_runs} subsequent runs: {total_exec_time / num_runs:.6f}s")

text

Execution took 0.098441s
Time for first run (compilation + execution): 0.098476s
Execution took 0.001954s
Execution took 0.002134s
Execution took 0.001834s
Execution took 0.002243s
Execution took 0.002433s
Average time for 5 subsequent runs: 0.002119s

Execution took 0.098441s
Time for first run (compilation + execution): 0.098476s
Execution took 0.001954s
Execution took 0.002134s
Execution took 0.001834s
Execution took 0.002243s
Execution took 0.002433s
Average time for 5 subsequent runs: 0.002119s

Python

@jax.jit
def jitted_computation_with_regular_print(a, b):
    c = a * b
    # INCORRECT for runtime debugging: Regular print executes during tracing.
    print(f"These messages will be printed only once, during tracing!")
    print(f"Regular print: Intermediate c: {c}")
    
    d = jnp.sin(c)
    print(f"Regular print: Shape of d: {d.shape}")

    e = jnp.sum(d)
    print(f"Regular print: Final sum e: {e}")
    return e

# Prepare some input arrays
input_a = jnp.array([1.0, 2.0, 3.0])
input_b = jnp.array([4.0, 5.0, 6.0])

print("\nDemonstrating regular print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with regular print():")
    result = jitted_computation_with_regular_print(input_a, input_b)
    result.block_until_ready() # Ensure computation finishes
    print(f"Result: {result}")
    print()

@jax.jit
def jitted_computation_with_regular_print(a, b):
    c = a * b
    # INCORRECT for runtime debugging: Regular print executes during tracing.
    print(f"These messages will be printed only once, during tracing!")
    print(f"Regular print: Intermediate c: {c}")
    
    d = jnp.sin(c)
    print(f"Regular print: Shape of d: {d.shape}")

    e = jnp.sum(d)
    print(f"Regular print: Final sum e: {e}")
    return e

# Prepare some input arrays
input_a = jnp.array([1.0, 2.0, 3.0])
input_b = jnp.array([4.0, 5.0, 6.0])

print("\nDemonstrating regular print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with regular print():")
    result = jitted_computation_with_regular_print(input_a, input_b)
    result.block_until_ready() # Ensure computation finishes
    print(f"Result: {result}")
    print()

text

Demonstrating regular print in a JIT-compiled function:
Call 1 with regular print():
These messages will be printed only once, during tracing!
Regular print: Intermediate c: Traced<float32[3]>with<DynamicJaxprTrace>
Regular print: Shape of d: (3,)
Regular print: Final sum e: Traced<float32[]>with<DynamicJaxprTrace>
Result: -2.0518109798431396

Call 2 with regular print():
Result: -2.0518109798431396

Call 3 with regular print():
Result: -2.0518109798431396

Demonstrating regular print in a JIT-compiled function:
Call 1 with regular print():
These messages will be printed only once, during tracing!
Regular print: Intermediate c: Traced<float32[3]>with<DynamicJaxprTrace>
Regular print: Shape of d: (3,)
Regular print: Final sum e: Traced<float32[]>with<DynamicJaxprTrace>
Result: -2.0518109798431396

Call 2 with regular print():
Result: -2.0518109798431396

Call 3 with regular print():
Result: -2.0518109798431396

Python

@jax.jit
def jitted_computation_for_debug(a, b):
    c = a * b
    # CORRECT: Using jax.debug.print to inspect 'c' during execution
    jax.debug.print("Inside JIT: Intermediate c: {c_val}", c_val=c)
    
    d = jnp.sin(c)
    # Using jax.debug.print to inspect shape of 'd'
    jax.debug.print("Inside JIT: Shape of d: {d_shape}", d_shape=d.shape)
    
    e = jnp.sum(d)
    jax.debug.print("Inside JIT: Final sum e: {e_val}", e_val=e)
    return e

print("\nDemonstrating jax.debug.print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with jax.debug.print():")
    debug_result = jitted_computation_for_debug(input_a, input_b)
    debug_result.block_until_ready() # Ensure computation and prints finish
    print(f"Result: {debug_result}")
    print()

@jax.jit
def jitted_computation_for_debug(a, b):
    c = a * b
    # CORRECT: Using jax.debug.print to inspect 'c' during execution
    jax.debug.print("Inside JIT: Intermediate c: {c_val}", c_val=c)
    
    d = jnp.sin(c)
    # Using jax.debug.print to inspect shape of 'd'
    jax.debug.print("Inside JIT: Shape of d: {d_shape}", d_shape=d.shape)
    
    e = jnp.sum(d)
    jax.debug.print("Inside JIT: Final sum e: {e_val}", e_val=e)
    return e

print("\nDemonstrating jax.debug.print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with jax.debug.print():")
    debug_result = jitted_computation_for_debug(input_a, input_b)
    debug_result.block_until_ready() # Ensure computation and prints finish
    print(f"Result: {debug_result}")
    print()

text

Demonstrating jax.debug.print in a JIT-compiled function:
Call 1 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 2 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 3 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Demonstrating jax.debug.print in a JIT-compiled function:
Call 1 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 2 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 3 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Python

# Brief example of using JAX profiler (conceptual)
import jax.profiler

jax.profiler.start_trace("/tmp/jax_profile_logs") # Start tracing, specify log directory

# # ... your complex JAX computations ...
result = timed_function(mat_x, mat_y) 

jax.profiler.stop_trace() # Stop tracing and save logs

# After running, you would visualize these logs using TensorBoard or other visualization tools

# Brief example of using JAX profiler (conceptual)
import jax.profiler

jax.profiler.start_trace("/tmp/jax_profile_logs") # Start tracing, specify log directory

# # ... your complex JAX computations ...
result = timed_function(mat_x, mat_y) 

jax.profiler.stop_trace() # Stop tracing and save logs

# After running, you would visualize these logs using TensorBoard or other visualization tools

Python

import jax
import jax.numpy as jnp
import time

# Create some test data for timing
key = jax.random.key(303)
size = 200
mat_x = jax.random.normal(key, (size, size))
key, _ = jax.random.split(key)
mat_y = jax.random.normal(key, (size, size))

Python

@jax.jit
def matrix_computation(x, y):
    for _ in range(5):  # Simulate some computational work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

Python

# Correct timing approach
start_time = time.perf_counter()
result = matrix_computation(mat_x, mat_y)
result.block_until_ready()  # Crucial: Wait for computation to complete on device
end_time = time.perf_counter()
accurate_time = end_time - start_time
print(f"Accurate execution time (first run, includes compilation): {accurate_time:.6f}s")

Python

import functools

def timer(func):
    """A decorator for timing function execution, with JAX support."""
    @functools.wraps(func) # Preserves original function metadata
    def wrapper(*args, **kwargs):
        start_time = time.perf_counter()
        result = func(*args, **kwargs)
        # For JAX arrays, block until computation is complete for accurate timing
        if hasattr(result, 'block_until_ready'):
            result.block_until_ready()
        end_time = time.perf_counter()
        run_time = end_time - start_time
        print(f"Execution took {run_time:.6f}s")
        return result, run_time # Return both result and runtime
    return wrapper

Python

@timer # Apply our custom timer
@jax.jit
def timed_function(x, y):
    for _ in range(5):  # Simulate some work
        x = jnp.dot(x, y) + 0.1
        y = jnp.dot(y, x) - 0.1
    return jnp.sum(x + y)

# First call (includes JIT compilation time)
_, compile_plus_exec_time = timed_function(mat_x, mat_y)
print(f"Time for first run (compilation + execution): {compile_plus_exec_time:.6f}s")

# Subsequent calls (execution only, JIT cache is used)
num_runs = 5
total_exec_time = 0
for _ in range(num_runs):
    _, run_time = timed_function(mat_x, mat_y)
    total_exec_time += run_time
print(f"Average time for {num_runs} subsequent runs: {total_exec_time / num_runs:.6f}s")

text

Execution took 0.098441s
Time for first run (compilation + execution): 0.098476s
Execution took 0.001954s
Execution took 0.002134s
Execution took 0.001834s
Execution took 0.002243s
Execution took 0.002433s
Average time for 5 subsequent runs: 0.002119s

Python

@jax.jit
def jitted_computation_with_regular_print(a, b):
    c = a * b
    # INCORRECT for runtime debugging: Regular print executes during tracing.
    print(f"These messages will be printed only once, during tracing!")
    print(f"Regular print: Intermediate c: {c}")
    
    d = jnp.sin(c)
    print(f"Regular print: Shape of d: {d.shape}")

    e = jnp.sum(d)
    print(f"Regular print: Final sum e: {e}")
    return e

# Prepare some input arrays
input_a = jnp.array([1.0, 2.0, 3.0])
input_b = jnp.array([4.0, 5.0, 6.0])

print("\nDemonstrating regular print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with regular print():")
    result = jitted_computation_with_regular_print(input_a, input_b)
    result.block_until_ready() # Ensure computation finishes
    print(f"Result: {result}")
    print()

text

Demonstrating regular print in a JIT-compiled function:
Call 1 with regular print():
These messages will be printed only once, during tracing!
Regular print: Intermediate c: Traced<float32[3]>with<DynamicJaxprTrace>
Regular print: Shape of d: (3,)
Regular print: Final sum e: Traced<float32[]>with<DynamicJaxprTrace>
Result: -2.0518109798431396

Call 2 with regular print():
Result: -2.0518109798431396

Call 3 with regular print():
Result: -2.0518109798431396

Python

@jax.jit
def jitted_computation_for_debug(a, b):
    c = a * b
    # CORRECT: Using jax.debug.print to inspect 'c' during execution
    jax.debug.print("Inside JIT: Intermediate c: {c_val}", c_val=c)
    
    d = jnp.sin(c)
    # Using jax.debug.print to inspect shape of 'd'
    jax.debug.print("Inside JIT: Shape of d: {d_shape}", d_shape=d.shape)
    
    e = jnp.sum(d)
    jax.debug.print("Inside JIT: Final sum e: {e_val}", e_val=e)
    return e

print("\nDemonstrating jax.debug.print in a JIT-compiled function:")
for i in range(1, 4):
    print(f"Call {i} with jax.debug.print():")
    debug_result = jitted_computation_for_debug(input_a, input_b)
    debug_result.block_until_ready() # Ensure computation and prints finish
    print(f"Result: {debug_result}")
    print()

text

Demonstrating jax.debug.print in a JIT-compiled function:
Call 1 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 2 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Call 3 with jax.debug.print():
Inside JIT: Shape of d: (Array(3, dtype=int32),)
Inside JIT: Intermediate c: [ 4. 10. 18.]
Inside JIT: Final sum e: -2.0518109798431396
Result: -2.0518109798431396

Python

# Brief example of using JAX profiler (conceptual)
import jax.profiler

jax.profiler.start_trace("/tmp/jax_profile_logs") # Start tracing, specify log directory

# # ... your complex JAX computations ...
result = timed_function(mat_x, mat_y) 

jax.profiler.stop_trace() # Stop tracing and save logs

# After running, you would visualize these logs using TensorBoard or other visualization tools