Skip to content

← JAX Performance & Distributed step 22 of 25

Medium Primitives

Compilation Cache Key

Jump to the task

Reuse the function you compiled

Suppose a training loop computes a sum of squares on each batch. The batch values change, but its shape and dtype often stay the same. A reusable JIT function lets JAX run the compiled calculation on each new batch without tracing the Python body again for every call.

import jax
import jax.numpy as jnp

@jax.jit
def sum_sq(x):
    return jnp.sum(x ** 2)

sum_sq(jnp.array([1., 2., 3.]))  # 14
sum_sq(jnp.array([2., 3., 4.]))  # 29; same shape and dtype
sum_sq(jnp.array([3., 4.]))      # 25; a different shape

The second call can reuse the first call's compiled calculation. It still computes a new answer: JIT caches executable code, not the previous numerical result. The two-element input needs a different specialization under ordinary JIT.

What changes the specialization?

Function identity, input shapes and dtypes, and explicit static argument values all matter. Other settings, including sharding and compilation options, can matter too; this is not a complete cache-key specification. Ordinary array values can change without retracing when their shape, dtype and other relevant properties stay the same.

Defining a new lambda or partial function on every call defeats function reuse:

def hot_call(x):
    fresh = jax.jit(lambda y: jnp.sum(y ** 2))
    return fresh(x)

Here the underlying Python function is new each time. Define the helper once and keep using it instead. Module scope is convenient; creating it once in a factory and retaining the returned function can work too. The important part is its lifetime, not its indentation. Retracing and backend compilation are distinct stages, so do not infer exact compilation counts or timing from this pattern alone.

Pass changing values as arguments

Python globals and captured values are read while tracing. Changing them later is not a reliable way to request a fresh specialization: a cached call can keep using the old value.

For example, if a jitted function computes x * scale with scale captured from Python, rebinding scale from 2 to 3 can leave a cached call multiplying by 2. Pass scale as an ordinary argument when it should vary numerically. Mark it static only when the calculation needs its concrete value during tracing, such as for Python control flow. Distinct static values can then require separate specializations.

The diagram below models one retained function under fixed settings. Its first two calls have different array values but the same float32 shape; the third changes the shape. It does not model function identity, captured values or persistent disk caches.

Loading visualization…

Problem

Implement jit_cached_call(x) using a module-level helper decorated with @jax.jit. Have the entrypoint call that existing helper and return jnp.sum(x ** 2).

X is a one-dimensional float32 JAX array. Return a scalar JAX array, not a Python float. For [3, 4], square to [9, 16] and sum to 25. An empty input sums to zero. Square each element before summing; squaring the sum would give 49 for [3, 4].

The output tests check numerical correctness. They do not prove that your helper is retained or that compilation was reused. Inspect the function placement as well; timing one call mixes setup and execution.