Skip to content

← JAX Performance & Distributed step 9 of 25

Medium Primitives

jit Retrace on Shape Change

Jump to the task

Why this matters

A jitted function can accept arrays of different lengths, but a new length normally needs another compiled version. Returning to a previous length can reuse the version already cached for it.

Tracing runs Python with placeholders to record JAX operations. Compilation turns that recorded computation into executable code. A cache hit avoids repeating this preparation; it still has to run the actual computation.

Worked mini-example

import jax
import jax.numpy as jnp

sum_jit = jax.jit(jnp.sum)
a = sum_jit(jnp.array([1.0, 2.0, 3.0]))       # 6, length 3
b = sum_jit(jnp.array([4.0, 3.0, 2.0, 1.0]))  # 10, length 4
c = sum_jit(jnp.array([7.0, 8.0, 9.0]))       # 24, length 3

The first two calls use different input shapes. The third can reuse the first version because only the values changed. This assumes the dtype and other compilation settings match and the cached version is still available.

Common pitfalls

  • Shape is only part of the cache key. Dtypes, static argument values, function identity and other JAX settings also matter.
  • A Python scalar and an explicitly typed scalar array can differ in JAX's weak-type metadata even when both have shape (). A scalar is not automatically a different shape from a scalar array.
  • Many shapes can increase compilation costs, but they do not automatically erase every benefit of JIT. Measure your workload.
  • Padding sequences to a few chosen lengths can reduce the number of versions. Mask padded entries so they do not change the answer.
  • JAX dispatch can be asynchronous. When measuring execution time, wait for the result with block_until_ready() and separate warm-up calls from later calls. There is no universal JIT latency.

Problem

Implement cache_independent_count(x_a, x_b):

  1. Define one jitted jnp.sum helper outside the entrypoint.
  2. Call that helper on each input.
  3. Add the two scalar results and return their sum.

Both inputs are one-dimensional JAX arrays; their lengths may differ. Despite the function's name, return the combined sum, not a cache entry count. The numeric tests do not measure compilation behavior.

For x_a = [1, 2] and x_b = [3], return 6. An empty input contributes zero.

The diagram shows lengths 8, 16, 32, then 8 again. With matching dtypes and settings, the last call can reuse the first version.

Loading visualization…