Avoiding Recompilation

JAX's jit decorator is a powerful tool for accelerating computations by compiling Python functions using XLA. This compilation process, however, isn't free. JAX first traces the Python function to generate an intermediate representation, the jaxpr, which captures the sequence of primitive operations. This jaxpr is then compiled by XLA into optimized code specific to the target hardware (CPU, GPU, TPU) and the shapes and types of the input arguments encountered during the trace.

The first time you call a @jit-decorated function with a particular combination of argument shapes and types, JAX performs this tracing and compilation, storing the resulting optimized executable in a cache. Subsequent calls with the same argument shapes and types can reuse the cached executable, making them significantly faster as they bypass the tracing and compilation overhead.

However, if you call the function with arguments whose characteristics differ from those seen before in a way that changes the jaxpr, JAX must re-trace and re-compile. This recompilation adds significant overhead, potentially negating the benefits of jit, especially if it happens frequently (e.g., inside a loop). Understanding and minimizing these recompilation events is important for achieving peak performance.

Why Recompilation Occurs

JAX's tracing mechanism works by substituting special tracer objects for the function's arguments. These tracers record the sequence of JAX primitive operations performed. The main point is that the trace depends on aspects of the input arguments:

1. Argument Shapes: If you call a function first with an array of shape (10, 5) and later with an array of shape (20, 5), the sequence of operations or their dimensions within the jaxpr might change. This triggers a re-trace and recompilation.
2. Argument Data Types (dtypes): Calling with float32 and then float64 arguments will also likely trigger recompilation, as different primitive implementations might be needed.
3. Argument Types: Mixing Python types like integers and floats for the same argument position across calls can cause recompilation. JAX needs to know whether an argument is a JAX array (traced) or a standard Python type (treated as static).
4. Value-Dependent Python Control Flow: If your function uses standard Python if statements or for loops whose conditions or iterations depend on the values of traced arguments (JAX arrays), the execution path taken during tracing can change. A different path means a different jaxpr, leading to recompilation.
5. Structure of Function Arguments: If you pass functions as arguments (common in higher-order functions), changing the function itself between calls necessitates recompilation.
6. Closures and Globals: If the JIT-compiled function closes over Python variables or accesses globals, and the values of these variables influence the computation graph's structure (e.g., determining the size of an intermediate array), changes to these values between calls can trigger recompilation.

Consider this simple example:

import jax
import jax.numpy as jnp
import time

@jax.jit
def example_function(x, size_param):
  # Python control flow depends on size_param's value
  if size_param > 5:
    return jnp.sum(x * 2.0)
  else:
    return jnp.sum(x + 1.0)

key = jax.random.PRNGKey(0)
data = jax.random.normal(key, (10,))

# First call: size_param = 3. Traces and compiles version 1.
print("First call...")
start_time = time.time()
result1 = example_function(data, 3)
result1.block_until_ready() # Ensure execution finishes
print(f"Result: {result1}, Time: {time.time() - start_time:.4f}s")

# Second call: size_param = 3 again. Uses cached version 1. Fast.
print("\nSecond call (cached)...")
start_time = time.time()
result2 = example_function(data, 3)
result2.block_until_ready()
print(f"Result: {result2}, Time: {time.time() - start_time:.4f}s")

# Third call: size_param = 7. Triggers re-trace and re-compile for version 2. Slow.
print("\nThird call (recompilation)...")
start_time = time.time()
result3 = example_function(data, 7)
result3.block_until_ready()
print(f"Result: {result3}, Time: {time.time() - start_time:.4f}s")

# Fourth call: size_param = 7 again. Uses cached version 2. Fast.
print("\nFourth call (cached)...")
start_time = time.time()
result4 = example_function(data, 7)
result4.block_until_ready()
print(f"Result: {result4}, Time: {time.time() - start_time:.4f}s")

You'll observe that the first and third calls are significantly slower due to the compilation overhead. The second and fourth calls reuse the cached executables and are much faster. The recompilation in the third call happened because the value of size_param changed the Python control flow path taken during tracing.

Strategies to Minimize Recompilation

While some recompilation is expected (e.g., when shapes genuinely change), frequent, unintended recompilation hampers performance. Here are important strategies to mitigate it:

1. Use Static Arguments (static_argnums / static_argnames)

When a function argument's value influences the computation graph structure (like shapes, hyperparameters determining layers, or values used in Python control flow), but the argument itself is not meant to be traced (e.g., it's a Python int or bool), you can mark it as "static".

JAX will treat static arguments as compile-time constants. It will trace and compile a specialized version of the function for each unique combination of values passed for these static arguments. This avoids recompilation when only the static arguments change, as JAX can look up the pre-compiled specialization.

You specify static arguments using the static_argnums (by index) or static_argnames (by name, often clearer) arguments to jax.jit:

import jax
import jax.numpy as jnp
import time

# Mark 'size_param' (index 1) as static
@jax.jit(static_argnums=(1,))
def example_function_static(x, size_param):
  print(f"Compiling for size_param = {size_param}") # See when compilation happens
  if size_param > 5:
    return jnp.sum(x * 2.0)
  else:
    return jnp.sum(x + 1.0)

key = jax.random.PRNGKey(0)
data = jax.random.normal(key, (10,))

# First call: size_param = 3. Compiles specialization 1.
print("First call...")
start_time = time.time()
result1 = example_function_static(data, 3)
result1.block_until_ready()
print(f"Result: {result1}, Time: {time.time() - start_time:.4f}s")

# Second call: size_param = 3 again. Uses cached specialization 1. Fast.
print("\nSecond call (cached)...")
start_time = time.time()
result2 = example_function_static(data, 3)
result2.block_until_ready()
print(f"Result: {result2}, Time: {time.time() - start_time:.4f}s")

# Third call: size_param = 7. Compiles specialization 2.
print("\nThird call (new static value)...")
start_time = time.time()
result3 = example_function_static(data, 7)
result3.block_until_ready()
print(f"Result: {result3}, Time: {time.time() - start_time:.4f}s")

# Fourth call: size_param = 7 again. Uses cached specialization 2. Fast.
print("\nFourth call (cached)...")
start_time = time.time()
result4 = example_function_static(data, 7)
result4.block_until_ready()
print(f"Result: {result4}, Time: {time.time() - start_time:.4f}s")

Now, the compilation happens only once for size_param=3 and once for size_param=7. Subsequent calls with these values reuse the appropriate cached specialization.

Caution: Use static arguments judiciously. If a static argument can take on many different values, you might end up compiling a large number of specializations, increasing compile time and memory usage for the compilation cache. It's best suited for arguments with a limited set of expected values that control the graph structure.

2. Ensure Consistent Input Types and Shapes

Whenever possible, try to call your JIT-compiled functions with arrays of consistent shapes and dtypes.

• Padding: If dealing with variable-length sequences, consider padding them to a fixed maximum length and using masking techniques (see Chapter 1) to ignore the padded elements.
• Type Casting: Explicitly cast inputs to the desired dtype before calling the JITted function if they might arrive with varying types.

3. Use JAX Control Flow Primitives

Replace Python if, for, and while loops that depend on traced values with their JAX counterparts: jax.lax.cond, jax.lax.scan, and jax.lax.while_loop. These primitives are integrated into the JAX tracing system. They embed the branching or looping logic within the compiled XLA graph, rather than creating different graphs based on Python execution paths. This avoids recompilation when the values controlling the flow change.

import jax
import jax.numpy as jnp
import time

# Use lax.cond instead of Python if
@jax.jit
def example_function_lax(x, size_param_val):
  # size_param_val must be a 0-dim array or Python scalar traceable by cond
  # Here we assume it's derived from data or passed appropriately
  pred = size_param_val > 5

  # Define functions for true and false branches
  def true_fun(operand):
      return jnp.sum(operand * 2.0)

  def false_fun(operand):
      return jnp.sum(operand + 1.0)

  # lax.cond selects which function to execute based on pred
  return jax.lax.cond(pred, true_fun, false_fun, x)

key = jax.random.PRNGKey(0)
data = jax.random.normal(key, (10,))

# Example: Pass a JAX scalar traceable by cond
# Usually, size_param_val would be computed from other JAX arrays
size_val_3 = jnp.array(3)
size_val_7 = jnp.array(7)

# First call: Compiles ONCE.
print("First call (lax.cond)...")
start_time = time.time()
result1 = example_function_lax(data, size_val_3)
result1.block_until_ready()
print(f"Result: {result1}, Time: {time.time() - start_time:.4f}s")

# Second call: Uses cached version. Fast. No recompilation.
print("\nSecond call (lax.cond, cached)...")
start_time = time.time()
result2 = example_function_lax(data, size_val_7)
result2.block_until_ready()
print(f"Result: {result2}, Time: {time.time() - start_time:.4f}s")

With lax.cond, only one compilation occurs, handling both branches within the XLA graph.

4. Manage Closures and Function Arguments Carefully

• Avoid closing over variables whose values change frequently and affect the graph structure. Pass such values as regular (potentially static) arguments instead.
• If passing functions as arguments (e.g., activation functions), ensure you pass the same function object each time. Using functools.partial can sometimes help create stable function objects if parameters need to be bound. Defining functions at the top level often avoids issues compared to defining them dynamically inside loops.

By actively identifying the causes of recompilation in your code using profiling (covered next) and applying these strategies, particularly static_argnums and JAX control flow primitives, you can significantly reduce compilation overhead and ensure your @jit-decorated functions run consistently fast after the initial compilation.