Skip to content

← JAX Performance & Distributed step 8 of 25

Medium Primitives

make_jaxpr Inspection

Jump to the task

Why this matters

A jaxpr records the JAX operations found while tracing a Python function. You can inspect it to check input types, intermediate shapes and the operations your code produces.

jax.make_jaxpr(f)(x) returns that representation instead of the numerical result of f(x). Python still runs during tracing, so ordinary Python side effects can occur.

Worked mini-example

import jax
import jax.numpy as jnp

x = jnp.ones(3, dtype=jnp.float32)
trace = jax.make_jaxpr(lambda x: jnp.sum(x ** 2))(x)
text = str(trace)
length = len(text)
print(text)

The trace describes squaring the elements and summing them. Shape annotations such as f32[3] describe a length-three float32 array; f32[] describes a float32 scalar. The exact names and formatting depend on the JAX version.

A jaxpr is not final machine code. Later compiler passes can change or combine operations, so neither equation count nor printed length tells you how many GPU kernels will run or how fast they will be.

Common pitfalls

  • Ordinary dynamic input values do not determine this trace, but shapes and dtypes can change it. Static arguments and captured constants can matter for other functions.
  • A literal in a trace is not automatically a bug. It may be an intentional constant written in the function.
  • Python branches resolve during tracing when their conditions are known. JAX control-flow operations can contain nested jaxprs.
  • Do not assume every inspection call reruns Python: tracing can be cached. Keep functions pure instead of relying on trace counts.
  • Different shapes can have the same printed length. Moving from a one-digit dimension to a two-digit dimension can change the length even when the operation sequence stays the same.

Problem

Implement jaxpr_string_length(x) for a one-dimensional float32 JAX array x.

  1. Trace exactly lambda x: jnp.sum(x ** 2) with jax.make_jaxpr using x as its example input.
  2. Convert the returned object to text with str(trace).
  3. Return jnp.float32(len(text)), a scalar array.

This is an exercise in inspecting a representation, not calculating the sum of squares. Do not hardcode a character count or substitute a mathematically equivalent function: it can produce different text.

The checker calculates the expected length on your runtime and also checks the scalar float32 return type. Formatting may change between JAX versions.