Static Arguments and Control Flow

Sometimes you need a function argument that affects the structure of the computation — like a flag that changes which branch to take, or an integer that determines the loop count. These are static arguments: values that JIT should treat as compile-time constants rather than traced values.

import jax
import jax.numpy as jnp
from functools import partial

# Problem: this fails because 'use_bias' is traced
@jax.jit
def linear_bad(x, w, b, use_bias):
    result = x @ w
    if use_bias:  # ConcretizationTypeError: can't use traced bool in if
        result = result + b
    return result

# Solution 1: static_argnums
@partial(jax.jit, static_argnums=(3,))
def linear_v1(x, w, b, use_bias):
    result = x @ w
    if use_bias:  # Now 'use_bias' is concrete — this works!
        result = result + b
    return result

# Solution 2: static_argnames (clearer)
@partial(jax.jit, static_argnames=('use_bias',))
def linear_v2(x, w, b, use_bias):
    result = x @ w
    if use_bias:
        result = result + b
    return result

x = jnp.ones((4, 3))
w = jnp.ones((3, 2))
b = jnp.ones(2)

result = linear_v2(x, w, b, use_bias=True)
print(result.shape)  # (4, 2)

Important: Changing a static argument triggers recompilation. Don't use static_argnums for values that change every call — that defeats the purpose of JIT.

JAX Control Flow Primitives

When you need value-dependent control flow inside JIT, use JAX's functional control flow primitives:

import jax
import jax.numpy as jnp

# jax.lax.cond: if-else for traced values
@jax.jit
def safe_divide(x, y):
    return jax.lax.cond(
        y != 0,
        lambda: x / y,         # true branch
        lambda: jnp.zeros_like(x)  # false branch
    )

# jax.lax.switch: multi-way branch (like switch/case)
@jax.jit
def activation(x, choice):
    return jax.lax.switch(choice, [
        lambda x: x,                          # 0: identity
        lambda x: jnp.maximum(x, 0),          # 1: relu
        lambda x: jnp.tanh(x),                # 2: tanh
        lambda x: jax.nn.sigmoid(x),           # 3: sigmoid
    ], x)

# jax.lax.fori_loop: for loop with traced bounds
@jax.jit
def power(x, n):
    return jax.lax.fori_loop(
        0, n,                    # start, stop (can be traced)
        lambda i, acc: acc * x,  # body: (iteration, carry) -> new_carry
        jnp.ones_like(x)        # initial carry
    )

# jax.lax.scan: the most powerful loop primitive
@jax.jit
def cumulative_product(arr):
    def step(carry, x):
        new_carry = carry * x
        return new_carry, new_carry  # (next carry, output)
    _, products = jax.lax.scan(step, jnp.array(1.0), arr)
    return products

print(cumulative_product(jnp.array([2.0, 3.0, 4.0])))  # [2. 6. 24.]