The price of a static argument

The clearest picture of what static_argnums does is the lowered program, before any of the caching argument comes up. Lower the same two-argument function twice, once plain and once with the second argument static, and the signatures do not match: the traced version takes tensor<4xf32> and tensor<f32>, the static version takes tensor<4xf32> alone and carries stablehlo.constant dense<2.000000e+00> in the body.

That is the whole mechanism, stated in the artifact rather than in prose. A static argument never becomes a parameter. Its value is read at trace time and written into the program, exactly the way a closed-over Python constant is, which is why chapter 03 describes it as baked in rather than passed in.

It also explains the shape of the cost. A constant in the program means the program is only valid for that constant, so a second value needs a second program, and the cache has no choice about it.

// jax.jit(scale).lower(jnp.ones(4), 2.0).as_text()
module @jit_scale attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<4xf32>, %arg1: tensor<f32>) -> (tensor<4xf32> {jax.result_info = ""}) {
    %0 = stablehlo.convert %arg1 : tensor<f32>
    %1 = stablehlo.broadcast_in_dim %0, dims = [] : (tensor<f32>) -> tensor<4xf32>
    %2 = stablehlo.multiply %arg0, %1 : tensor<4xf32>
    return %2 : tensor<4xf32>
  }
}

// jax.jit(scale, static_argnums=1).lower(jnp.ones(4), 2.0).as_text()
module @jit_scale attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<4xf32>) -> (tensor<4xf32> {jax.result_info = ""}) {
    %cst = stablehlo.constant dense<2.000000e+00> : tensor<f32>
    %0 = stablehlo.broadcast_in_dim %cst, dims = [] : (tensor<f32>) -> tensor<4xf32>
    %1 = stablehlo.multiply %arg0, %0 : tensor<4xf32>
    return %1 : tensor<4xf32>
  }
}

A folded constant sounds like free speed, so it is worth asking what it actually bought. Take a 2048 by 2048 elementwise scale-and-add, compile it both ways, and time the steady state. On this machine the traced version ran at 5.33 milliseconds and the static version at 5.40, which is to say the two are the same number wearing different error bars. XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → broadcasts a scalar parameter about as cheaply as it broadcasts a scalar constant, and there was nothing else in the program for the constant to unlock.

The other side of the ledger is not within noise at all. The first call at a new static value cost 139.22 milliseconds against a 5.33 millisecond steady call, so one new value costs roughly twenty-six calls of debt and returns nothing measurable on this program. Converting a compile into steady steps is chapter 11's habit, taught on a six-layer MLP in its lesson on what the clock caught; it is borrowed here to price a single keyword.

So the decision is not about whether specialization is good in general. It is arithmetic over two numbers: how many distinct values this argument takes across a run, and what one compile costs at this shape. LAB·J2 counts the values for you with the compile log. What this lesson adds is the other factor, measured for the case where the answer comes out against static: a scalar that changes no shape.

import statistics as st
import time

import jax
import jax.numpy as jnp

x = jnp.ones((2048, 2048))

def scale_a(x, s):
    return x * s + 1.0

def scale_b(x, s):
    return x * s + 1.0

traced = jax.jit(scale_a)
static = jax.jit(scale_b, static_argnums=1)

def median_ms(call, n=9):
    ts = []
    for _ in range(n):
        t0 = time.perf_counter()
        call().block_until_ready()
        ts.append(1e3 * (time.perf_counter() - t0))
    return round(st.median(ts), 2)

traced(x, 2.0).block_until_ready()
static(x, 2.0).block_until_ready()
print("s traced, steady state", median_ms(lambda: traced(x, 2.0)), "ms")
print("s static, steady state", median_ms(lambda: static(x, 2.0)), "ms")

t0 = time.perf_counter()
static(x, 3.0).block_until_ready()
print("s static, one new value", round(1e3 * (time.perf_counter() - t0), 2), "ms")
print("entries: traced", traced._cache_size(), "static", static._cache_size())

# s traced, steady state 5.33 ms
# s static, steady state 5.4 ms
# s static, one new value 139.22 ms
# entries: traced 1 static 2

A static value has to hash. The museum files the error you get when it does not, under a list handed to static_argnames, so that half of the contract already has a home. The other half fails silently instead of raising, which is why it is worth a section here.

Two static values land on the same cache entry when they are equal, and a plain Python object without __eq__ compares by identity. So a small config object, constructed fresh at each call site with the same contents, is a different key every time. Three equal-looking configs produce three cache entries and three compiles, and nothing anywhere prints a warning.

Give the class __hash__ and __eq__ by value and the same three calls collapse to one entry. That is the fix for the config-object pattern generally: a frozen dataclass, a NamedTuple, or an explicit pair of dunders. The compiler is comparing your keys with ==, so whatever == means for your type is what the cache means by the same program.

from functools import partial

import jax
import jax.numpy as jnp

x = jnp.ones(4)

class Cfg:
    def __init__(self, n):
        self.n = n

class KeyedCfg:
    def __init__(self, n):
        self.n = n
    def __hash__(self):
        return hash(self.n)
    def __eq__(self, other):
        return isinstance(other, KeyedCfg) and self.n == other.n

@partial(jax.jit, static_argnums=1)
def loose(x, cfg):
    return x * cfg.n

@partial(jax.jit, static_argnums=1)
def keyed(x, cfg):
    return x * cfg.n

for _ in range(3):
    loose(x, Cfg(2))
    keyed(x, KeyedCfg(2))

print("three equal Cfg objects     ", loose._cache_size())
print("three equal KeyedCfg objects", keyed._cache_size())

# three equal Cfg objects      3
# three equal KeyedCfg objects 1

None of the above says avoid static_argnums. It says spend it where a traced argument cannot do the job at all, which is any argument the output shape depends on. A repeat count, a number of heads, a window size, a boolean that selects between two differently shaped returns: each of those has to be known while the trace is running, because the trace is where shapes get fixed.

For those arguments the cost analysis inverts. The cardinality is small and bounded by the model rather than by the data, so a handful of compiles buys programs that could not otherwise exist. A per-example sequence length, by contrast, is bounded by the data, and marking it static is the churn bug the chapter warns about.

There is a third option that is easy to forget between the two. An argument whose value is genuinely fixed for the life of the program does not need to be an argument: close over it, and it is baked in at trace time for free, with no cache key component to keep track of and no way for a caller to accidentally vary it.