Speeding Up with jax.jit: Just-In-Time Compilation

Introduction

Welcome back to the fourth lesson of "JAX Fundamentals: NumPy Power-Up"! We've made excellent progress together: you've mastered JAX arrays and their immutable nature, understood the critical importance of pure functions, and learned how to compute gradients automatically with jax.grad . Now, we're ready to unlock a significant performance boost for our NumPy-like code with one of JAX's most transformative features: Just-In-Time (JIT) compilation . Today, we'll discover how jax.jit can dramatically accelerate our numerical computations by compiling Python functions into highly optimized machine code. As you may recall, pure functions enable JAX's powerful transformations, and JIT compilation is perhaps the most immediately rewarding of these. We'll learn how to apply JIT compilation, understand its initial overhead, measure performance improvements, and handle common challenges like control flow within compiled functions. By the end of this lesson, you'll understand when and how to use jax.jit effectively, and you'll have the tools to achieve significant speedups in your numerical computations, truly "powering up" your JAX skills.

Understanding Just-In-Time Compilation

Before we start speeding up our code, let's understand what Just-In-Time compilation actually does and why it's so powerful for numerical computing. When we write Python code using JAX operations, we're essentially describing a sequence of mathematical operations. However, Python itself is an interpreted language, meaning each operation is typically executed one at a time, which can introduce overhead. JIT compilation takes a different approach: instead of executing operations one by one, JAX analyzes the entire sequence of operations within a function and translates them into highly optimized machine code using XLA (Accelerated Linear Algebra) . XLA is Google's domain-specific compiler for linear algebra that can produce extremely efficient code for CPUs, GPUs, and TPUs. The "Just-In-Time" aspect means this compilation happens the first time we call a function with specific input shapes and types, not when we define it. This allows JAX to tailor optimizations. The compiled version is then cached and reused for subsequent calls with compatible inputs, giving us the flexibility of Python with the performance of compiled code. Think of it like this: if you're assembling furniture, you could read the instructions and find each tool as you need it (like an interpreter). Or, you could study the entire manual first, lay out all the tools in optimal order, and create an efficient assembly line (like a JIT compiler). Setting up this optimized workflow takes time initially (compilation), but once established, you can assemble identical furniture much faster than the step-by-step approach. The key requirement for JIT compilation is that our functions must be pure — exactly what we learned about in our second lesson. Pure functions, with no side effects and deterministic outputs, allow JAX to safely analyze and optimize them.

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

# Define the SELU activation function
def selu(x, alpha=1.67, lambda_=1.05):
  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

# Create a large input array to see meaningful timing differences
size = 10000000
vector = jnp.arange(size, dtype=jnp.float32) # Using float32 for consistency
print(f"Input size: {vector.shape}")

text

Input size: (10000000,)

Python

# Running without JIT - warm-up run first
start_time_first = time.time()
_ = selu(vector).block_until_ready() # Ensure computation completes
end_time_first = time.time()
print(f"First non-JIT run: {(end_time_first - start_time_first):.4f} seconds")

# Time multiple runs without JIT for an average
start_time_nojit_runs = time.time()
for _ in range(5):
    result_nojit = selu(vector)
    result_nojit.block_until_ready() # Ensure each computation completes
end_time_nojit_runs = time.time()
print(f"Time for 5 runs (no JIT): {(end_time_nojit_runs - start_time_nojit_runs):.4f} seconds")

Python

# Apply JIT compilation to the selu function
jit_selu = jax.jit(selu)

# First call to the JIT-compiled function includes compilation time
start_time_first_jit = time.time()
result_jit_first = jit_selu(vector)
result_jit_first.block_until_ready() # Wait for compilation and execution
end_time_first_jit = time.time()
print(f"First JIT run (compilation + execution): {(end_time_first_jit - start_time_first_jit):.4f} seconds")

# Subsequent calls use the cached, compiled version
start_time_subsequent_jit = time.time()
for _ in range(5):
    result_jit_subsequent = jit_selu(vector)
    result_jit_subsequent.block_until_ready() # Wait for execution
end_time_subsequent_jit = time.time()
print(f"Time for 5 subsequent JIT runs: {(end_time_subsequent_jit - start_time_subsequent_jit):.4f} seconds")

# Verify that the results are the same
print(f"Are results the same? {jnp.allclose(result_nojit, result_jit_subsequent)}")

text

First non-JIT run: 0.1505 seconds
Time for 5 runs (no JIT): 0.5007 seconds
First JIT run (compilation + execution): 0.0313 seconds
Time for 5 subsequent JIT runs: 0.0761 seconds
Are results the same? True

Python

@jax.jit
def selu_decorated(x, alpha=1.67, lambda_=1.05):
  return lambda_ * jnp.where(x > 0, x, alpha * jnp.exp(x) - alpha)

# Now selu_decorated is automatically JIT-compiled

Python

def conditional_computation(x, use_add: bool):
    """A function with Python control flow based on an argument."""
    if use_add: # Python if statement
        return x + 5.0
    else:
        return x * 2.0

input_val = jnp.array(10.0)

# Try to JIT compile this function directly
print("Trying JIT without static_argnames (this will fail):")
try:
    broken_jit_conditional = jax.jit(conditional_computation)
    result = broken_jit_conditional(input_val, use_add=True) # Pass boolean argument
    result.block_until_ready()
except Exception as e:
    print(f"Error: {type(e).__name__}: {str(e)}")

text

Trying JIT without static_argnames (this will fail):
Error: TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function conditional_computation at main.py:9 for jit. This concrete value was not available in Python because it depends on the value of the argument use_add.

Python

# Correct way: mark 'use_add' as a static argument
jit_conditional_computation = jax.jit(conditional_computation, static_argnames='use_add')

# Equivalent alternative, using static_argnums to specify index:
# jit_conditional_computation = jax.jit(conditional_computation, static_argnums=1)


# Call with use_add=True (compiles a version for use_add=True)
print("Calling with use_add=True:")
res_add = jit_conditional_computation(input_val, use_add=True)
res_add.block_until_ready()
print(f"Result: {res_add}")

# Call with use_add=False (compiles a version for use_add=False)
print("Calling with use_add=False:")
res_mul = jit_conditional_computation(input_val, use_add=False)
res_mul.block_until_ready()
print(f"Result: {res_mul}")

# Call again with use_add=True (uses the cached version for use_add=True)
print("Calling with use_add=True again (cached):")
res_add_again = jit_conditional_computation(input_val, use_add=True)
res_add_again.block_until_ready()
print(f"Result: {res_add_again}")

text

Calling with use_add=True:
Result: 15.0
Calling with use_add=False:
Result: 20.0
Calling with use_add=True again (cached):
Result: 15.0

Introduction

Understanding Just-In-Time Compilation

Join the 1M+ learners on CodeSignal

Your First JIT-Compiled Function

Measuring Performance: Before and After JIT

The Magic of JIT Compilation

jax.jit as Decorator

Problem: Python Control Flow in JIT

Solution: Static Arguments

Best Practices and Common Pitfalls

Conclusion and Next Steps