Skip to content

← JAX Profiling & Debugging step 12 of 12

Medium Primitives

Pretty-Printed jaxpr Length

Jump to the task

Inspect the text of a traced computation

jax.make_jaxpr records the JAX operations in a function for the supplied input shape and dtype. Its printed representation is useful for reading a trace. It is not machine code, and its character count is not a measure of runtime, memory use, or the number of fused device kernels.

This exercise practices obtaining that text and measuring its length. The exact formatting is a debugging representation, not a stable file format. Compute the answer from the trace on your installed JAX version.

Worked example

import jax
import jax.numpy as jnp

x = jnp.ones(3, dtype=jnp.float32)
closed = jax.make_jaxpr(lambda x: jnp.sum(x ** 2))(x)
text = closed.jaxpr.pretty_print()
print(text)
length = len(text)
print(length)  # depends on the installed JAX version and formatting

Read the trace as a sequence of bindings: square the input, then reduce those values to a sum. Type annotations describe array shapes and dtypes. The final in (...) lists the returned values. The input's actual values are not substituted into this sum-of-squares trace.

Problem

Implement pp_jaxpr_len(x):

  1. Trace exactly lambda x: jnp.sum(x ** 2) with jax.make_jaxpr on x.
  2. Access .jaxpr on the returned trace object.
  3. Call .pretty_print() with its default options.
  4. Return jnp.float32(len(text)) for the resulting string.

x is a one-dimensional float32 JAX array. Return a scalar JAX value with dtype float32. Do not return the text, count its lines, remove whitespace, or hard-code a length from the example.

The exercise checker measures the same representation using your runtime's JAX installation, then compares your result with that length. It also checks the scalar shape and float32 dtype. This avoids tying correctness to one release's whitespace or printed primitive parameters.

What can change the length

Two traces can perform the same mathematical calculation and print differently. Primitive names, parameters and formatting can change across JAX releases. Different Python formulations can also produce different traces, even when their numerical results agree. Trace the function given in the task rather than replacing it with an equivalent expression.

Shape annotations matter too. Changing a dimension from 3 to 5 may leave the character count unchanged, but changing it from 9 to 10 adds digits. A similar operation sequence does not guarantee equal printed lengths.

Do not infer a general formatting difference between str(closed) and closed.jaxpr.pretty_print(): their output can coincide, and the object wrappers have changed across JAX releases. Use the exact call requested here. When saving a trace for debugging, include the JAX version, input shape and dtype so another person can interpret it.