Static vs Traced Values

As we saw in the previous section, jax.jit works its magic by tracing your Python function. It executes the function once with placeholder objects, called tracers, representing the inputs. These tracers record all the operations performed on them, building an intermediate representation (a Jaxpr) which is then compiled by XLA into optimized machine code.

This tracing mechanism leads to an important distinction: some values involved in your function are traced, while others must be treated as static. Understanding this difference is significant for correctly and efficiently using jit.

Traced Values (Tracers)

Most often, the inputs to your JIT-compiled functions will be JAX arrays (or structures containing JAX arrays). When jit traces the function, these array inputs are replaced by tracer objects.

• What they are: Tracers are symbolic placeholders that stand in for the actual arrays during compilation. They primarily capture the shape and data type (dtype) of the expected input arrays.
• How they behave: When a tracer encounters a JAX operation (like jnp.dot, jnp.sin, etc.), it doesn't perform the calculation immediately. Instead, it records that operation as part of the computation graph (Jaxpr). The actual numerical computations happen later, after compilation, when you call the compiled function with concrete data.
• Value Availability: The concrete numerical values of traced inputs are not available during the tracing process itself. Only their shape and dtype are known.

Think of it like drawing a blueprint for a calculation. The blueprint specifies the dimensions and types of materials (shape and dtype) and the steps to follow (operations), but you don't need the actual physical materials (concrete values) until you start building (executing the compiled code).

Static Values

Static values, in contrast, are values that are known and fixed at compile time (during tracing). They are treated as constants within the compiled code.

• What they are: These are typically Python primitives like integers, booleans, strings, or function arguments explicitly marked as static. They can also include things like the shape tuples derived from traced inputs.
• How they behave: Operations or control flow constructs that depend on static values are executed during tracing. The result of these operations becomes part of the compiled function's structure. For instance, if an if statement's condition evaluates to a static True, only the code within that if block will be included in the compiled Jaxpr.
• Value Availability: The concrete values of static inputs are available during tracing.

Why Does This Distinction Matter?

The core issue arises when standard Python control flow (like if statements or for loops) depends on the value of an input that JAX treats as a traced value by default.

Consider this simple function:

import jax
import jax.numpy as jnp

def conditional_double(x, threshold):
  # Python if statement depends on the value of x
  if x > threshold:
    return x * 2
  else:
    return x

Let's try to JIT-compile it:

jitted_conditional_double = jax.jit(conditional_double)

# This will likely raise an error!
try:
  result = jitted_conditional_double(jnp.array(5.0), threshold=0.0)
  print(result)
except Exception as e:
  print(f"Error: {e}")

You'll encounter an error similar to this: ConcretizationTypeError: Abstract tracer value encountered where concrete value is expected....

Why the error? During tracing, x is replaced by a tracer. The Python if statement attempts to evaluate tracer > threshold. But the tracer doesn't have a concrete numerical value; it only knows its shape and dtype. It cannot produce a single boolean True or False needed for the Python if at compile time. Python needs a concrete boolean value right then and there to decide which branch to trace, but JAX only has a placeholder.

Specifying Static Arguments: static_argnums and static_argnames

To handle situations where control flow or other function logic must depend on the concrete value of an argument during tracing, JAX provides ways to mark arguments as static.

You can tell jit to treat specific arguments as static using the static_argnums or static_argnames arguments:

• static_argnums: Provide a tuple of integers specifying the positional indices of arguments that should be static.
• static_argnames: Provide a tuple of strings specifying the names of arguments (positional or keyword) that should be static.

Let's fix our previous example. Suppose we know the threshold will likely be constant for many calls, or we need the if statement to work as standard Python logic. We can mark threshold as static:

# Using static_argnums (threshold is the 1st argument, index 1)
jitted_conditional_double_nums = jax.jit(conditional_double, static_argnums=(1,))

# Using static_argnames
jitted_conditional_double_names = jax.jit(conditional_double, static_argnames=('threshold',))

# Now these work:
result1 = jitted_conditional_double_nums(jnp.array(5.0), threshold=0.0)
print(f"Result (static_argnums): {result1}") # Output: Result (static_argnums): 10.0

result2 = jitted_conditional_double_names(jnp.array(-2.0), threshold=0.0)
print(f"Result (static_argnames): {result2}") # Output: Result (static_argnames): -2.0

Now, when jit traces the function, it knows that threshold is static. It substitutes the actual value provided for threshold (e.g., 0.0) during the trace. The Python if statement can then evaluate tracer > 0.0. While this still involves a tracer, JAX can sometimes handle comparisons involving constants and tracers by embedding conditional logic into the compiled code using specialized primitives (like lax.cond). However, more complex Python logic depending on traced values will still fail. Making the value controlling the Python logic static ensures the Python interpreter can execute the control flow during tracing.

The Recompilation Trade-off

Marking arguments as static allows more standard Python constructs within JIT-compiled functions, but it comes with a significant performance consideration: recompilation.

• If you call a JIT-compiled function with different values for its traced arguments (but the same shapes and dtypes), JAX reuses the already compiled code. This is fast.
• If you call a JIT-compiled function with different values for its static arguments, JAX must re-trace and re-compile the function for those new static values. This compilation overhead can negate the benefits of jit if static arguments change frequently.

# threshold is static here
jitted_func = jax.jit(conditional_double, static_argnames=('threshold',))

print("First call (threshold=0.0):")
_ = jitted_func(jnp.array(5.0), threshold=0.0) # Compiles for threshold=0.0

print("Second call (threshold=0.0):")
_ = jitted_func(jnp.array(1.0), threshold=0.0) # Reuses compiled code

print("Third call (threshold=10.0):")
_ = jitted_func(jnp.array(5.0), threshold=10.0) # *** Recompiles for threshold=10.0 ***

print("Fourth call (threshold=10.0):")
_ = jitted_func(jnp.array(1.0), threshold=10.0) # Reuses compiled code for threshold=10.0

JAX caches compiled functions based on the identity of the Python function object and the static argument values (along with input shapes/dtypes).

When to Use Static Arguments

You generally need static arguments when:

1. Python Control Flow Depends on Argument Value: Standard Python if, for, while loops whose conditions or iterations depend directly on an argument's concrete value (not just its shape).
2. Function Behavior Parameterization: An argument determines which function to call or specifies structural aspects of the computation (e.g., number of layers in a neural network, type of activation function).
3. Array Indexing with Variables: Indexing arrays using values derived from arguments might require those arguments to be static if the indices need to be constant for compilation.

General Guideline: Prefer traced arguments (the default) for numerical data (JAX arrays) that changes often. Use static arguments sparingly for values that control the structure of the computation or are required by Python's runtime logic during tracing, and which do not change frequently between calls.

If you find yourself needing dynamic control flow based on traced values, explore JAX's structured control flow primitives like lax.cond, lax.scan, and lax.switch, which are designed to be traceable and compilable. We will touch upon these later, but they offer a way to express dynamic computation graphs compatible with JIT.

Understanding the difference between static and traced values is essential for debugging jit issues and optimizing performance by minimizing recompilation. By carefully considering which arguments need to be static, you can effectively accelerate your JAX code.