Skip to content

← JAX Conceptual Deep-Dives step 15 of 15

Medium Primitives

jaxpr with jit

Jump to the task

Inspect a trace, then run the function

jax.make_jaxpr(f)(x) shows a JAX-level representation of a function for an input shape and dtype. jax.jit(f) creates a callable that can compile and execute that computation. These are related tools, but inspecting a trace is not the same as running the compiled calculation.

When tracing a JIT-wrapped function, the outer trace may contain a call with a nested jaxpr for the function body. The call's name and presentation can change across releases. Read the nested computation rather than relying on a particular printed primitive name.

Worked example

import jax
import jax.numpy as jnp

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

x = jnp.array([1.0, 2.0, 3.0], dtype=jnp.float32)
trace = jax.make_jaxpr(sum_squares)(x)
print(trace)                 # inspect the traced structure
result = sum_squares(x)      # execute the calculation
# result is 14.0

The result is 1 + 4 + 9 = 14. The trace describes how to calculate it; it is not the scalar result itself.

Problem

Implement jaxpr_inside_jit(x) to return the sum of squares using a JIT-compiled helper. Define the helper once at module scope with @jax.jit, then call it from jaxpr_inside_jit.

  • x is a one-dimensional float32 JAX array.
  • Return a scalar equal to jnp.sum(x ** 2).
  • An empty input returns zero.
  • You may inspect the helper with jax.make_jaxpr while learning, but do not return that trace from this function.

The numerical tests check the returned sum. They do not prove that an implementation used JIT or that it runs faster.

Keep the compiled helper reusable

Defining and decorating a fresh helper inside every call creates a new Python function each time and can defeat compilation reuse. Keep the helper outside the wrapper so later calls can reuse it.

For a reusable function, array shapes and dtypes affect specialization, as do static arguments and other compilation settings. Changing only the array values normally does not require a new specialization. A displayed jaxpr is not a record of cache hits or a diagnosis of what caused recompilation. Tracing itself may be cached; do not assume every make_jaxpr invocation executes the Python body again.

Look at the appropriate stage

A jaxpr can contain high-level JAX primitives such as convolution. It is not optimized device code, and equation counts do not reveal kernel fusion. For a lower-level view, use sum_squares.lower(x).compiler_ir(); for text from a compiled executable, use sum_squares.lower(x).compile().as_text(). The available details depend on the backend and JAX version. These are inspection tools, not portable file formats or timing measurements.

The diagram below illustrates calls to the same reusable function with one unchanged shape followed by a changed shape. The labels x32 and x64 refer to array lengths, not floating-point dtypes.

Loading visualization…