Two counters, one cache

Chapter 02 established that a Python side effect inside a jitted function fires once, at trace time, and never again. Read that as a measuring instrument rather than a warning and you get the cheapest recompile detector there is: a global counter incremented in the function body ticks exactly once per trace, so its value is the number of times JAX has walked your Python.

The second number comes from the jit object itself. f._cache_size() reports how many entries the dispatch cache holds for the function f wraps. The leading underscore is a real caveat and worth saying out loud: it is not part of the promised API, and a future release may rename it. As an instrument in a lesson it is fine, and nothing else in the public surface answers the same question.

Run the two together and the first three lines are unsurprising. The same argument twice moves neither counter, which is the cache doing its job.

import jax
import jax.numpy as jnp
import numpy as np

traces = 0

@jax.jit
def f(x, s):
    global traces
    traces += 1          # a trace-time side effect: one tick per trace
    return x * s

x = jnp.ones(4)

def call(label, s):
    f(x, s)
    print(f"{label:<16} traces={traces} entries={f._cache_size()}")

call("jnp float32", jnp.float32(2.0))
call("the same again", jnp.float32(2.0))
call("numpy float32", np.float32(2.0))
call("python float", 2.0)
call("python int", 2)

# jnp float32      traces=1 entries=1
# the same again   traces=1 entries=1
# numpy float32    traces=1 entries=2
# python float     traces=2 entries=3
# python int       traces=3 entries=4

Line three is where the instrument earns its place. Swapping jnp.float32(2.0) for np.float32(2.0) adds a cache entry and does not add a trace. Nothing about the program changed, because both scalars abstract to the same shape and dtype, so JAX reused the jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr → it already had. What it did not reuse was the fast dispatch path, which is keyed more finely than the trace is, and a NumPy scalar reaches it as a different kind of Python object than a JAX array does.

Lines four and five move both counters, and chapter 01's weak scalars are why. A Python 2.0 abstracts to float32[] with weak_type=True, a Python 2 to int32[] with weak_type=True, and each of those is a different abstract value from the strong float32[] the first three calls produced. A different abstract value is a different program, so it is a real retrace.

The practical reading is a rule about argument hygiene rather than about dtypes. Passing a hyperparameter as a bare Python number, and sometimes as an array, and sometimes as a NumPy scalar, spends cache entries on a value that never changed. Pick one form at the call site and the counters flatten out.

Two numbers that move apart are telling you which layer noticed.

import jax.numpy as jnp
import numpy as np

for label, v in [("jnp.float32(2.0)", jnp.float32(2.0)),
                 ("np.float32(2.0)", np.float32(2.0)),
                 ("2.0", 2.0),
                 ("2", 2)]:
    a = jnp.asarray(v)
    print(f"{label:<17} {a.dtype} weak_type={a.weak_type}")

# jnp.float32(2.0)  float32 weak_type=False
# np.float32(2.0)   float32 weak_type=False
# 2.0               float32 weak_type=True
# 2                 int32 weak_type=True

Pytree structure is the part of the key people most often guess wrong about, in both directions. Handing the same two arrays as a dict whose keys are written in a different order costs nothing, because flattening sorts dict keys before it builds the treedef, so {"b": x, "a": x} and {"a": x, "b": x} produce the same PyTreeDef and land on the same entry.

Handing them as a list instead of a tuple costs a full retrace. PyTreeDef((*, *)) and PyTreeDef([*, *]) are different structures, and JAX makes no attempt to treat one container as the other. The counter proves it in one line, and printing the treedef shows you exactly what the cache compared.

That asymmetry is worth carrying into real code. A data pipeline that yields a list on one step and a tuple on the next, or a config that arrives as a dict here and a NamedTuple there, is a recompile per shape of container, and none of it shows up as a change in shapes or dtypes.

import jax
import jax.numpy as jnp

traces = 0

@jax.jit
def g(t):
    global traces
    traces += 1
    if isinstance(t, dict):
        return t["a"] + t["b"]
    return t[0] + t[1]

x = jnp.ones(3)
g({"a": x, "b": x}); print("dict a,b ", traces, g._cache_size())
g({"b": x, "a": x}); print("dict b,a ", traces, g._cache_size())
g((x, x));           print("tuple    ", traces, g._cache_size())
g([x, x]);           print("list     ", traces, g._cache_size())

print(jax.tree_util.tree_structure({"b": x, "a": x}))
print(jax.tree_util.tree_structure((x, x)))
print(jax.tree_util.tree_structure([x, x]))

# dict a,b  1 1
# dict b,a  1 1
# tuple     2 2
# list      3 3
# PyTreeDef({'a': *, 'b': *})
# PyTreeDef((*, *))
# PyTreeDef([*, *])

The advice you usually hear is to hoist jax.jit(f) out of the loop, and the counters say that advice is aimed at the wrong object. Build a fresh jax.jit wrapper around the same named function on every iteration and the calls stay at cache-hit speed, because the key holds the function jit wraps, not the wrapper. jax.jit(step)._cache_size() still reports one entry after all of it.

Build a fresh function on every iteration and the cost shows up on the first call. A lambda written inside a loop, or an inner function returned by a factory, is a new function object each time round, so it is a new key, so it traces and compiles from scratch. In the run below that is 0.4 milliseconds against 83.3, on a computation small enough that the arithmetic itself is a rounding error, and on a busier machine the gap widens rather than closes.

The factory pattern is the version that shows up in real training code, because it reads as configuration rather than as churn. make_step(lr) called once per epoch to bake in a new learning rate hands jit a new function every epoch, and every one of them compiles. Building it once and passing the learning rate as an argument costs one compile in total.

import statistics as st
import time

import jax
import jax.numpy as jnp

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

def step(v):
    return v * 2.0 + 1.0

def make_step(scale):
    def inner(v):
        return v * scale + 1.0
    return inner

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

jax.jit(step)(x).block_until_ready()          # compile once, for real
jax.jit(make_step(2.0))(x).block_until_ready()

print("a new wrapper over one function", median_ms(lambda: jax.jit(step)(x).block_until_ready()), "ms")
print("a new closure from a factory  ", median_ms(lambda: jax.jit(make_step(2.0))(x).block_until_ready()), "ms")
print("entries under step:", jax.jit(step)._cache_size())

# a new wrapper over one function 0.4 ms
# a new closure from a factory   83.3 ms
# entries under step: 1