The residuals you refuse to keep

Reverse mode holds values from the forward pass until the backward pass consumes them, and print_saved_residuals prints that list without executing the program. Four tanh layers over a 256 by 512 activation save nine arrays: two per layer, plus one from the final sum.

Each line names the shape, the primitive that produced the value, and the source position it came from, so the list is a map back into your own code rather than a total. Nine arrays of f32[256,512] is 4.5 MiB at this size, and the same function at sixteen layers saves thirty-three of them.

Wrap the same function in jax.checkpoint and the list collapses to the two arguments. No activation survives the forward pass; the backward pass will recompute what it needs. That is the change remat makes, printed rather than argued.

One import detail on this version: jax.ad_checkpoint is not reachable as an attribute of jax, and reaching for it raises AttributeError: module 'jax' has no attribute 'ad_checkpoint'. Did you mean: 'checkpoint'?. Import the name from the submodule, as below.

import jax
import jax.numpy as jnp
from jax.ad_checkpoint import print_saved_residuals

W = jax.random.normal(jax.random.key(0), (512, 512)) * 0.05
x = jnp.ones((256, 512))

def stack(W, x):
    h = x
    for _ in range(4):
        h = jnp.tanh(h @ W)
    return jnp.sum(h ** 2)

print_saved_residuals(stack, W, x)
print("--- under jax.checkpoint")
print_saved_residuals(jax.checkpoint(stack), W, x)

# f32[512,512] from the argument W
# f32[256,512] from the argument x
# f32[256,512] output of tanh from <string>:11:12 (stack)
# f32[256,512] output of sub from <string>:11:12 (stack)
# f32[256,512] output of tanh from <string>:11:12 (stack)
# f32[256,512] output of sub from <string>:11:12 (stack)
# f32[256,512] output of tanh from <string>:11:12 (stack)
# f32[256,512] output of sub from <string>:11:12 (stack)
# f32[256,512] output of tanh from <string>:11:12 (stack)
# f32[256,512] output of sub from <string>:11:12 (stack)
# f32[256,512] output of mul from <string>:12:19 (stack)
# --- under jax.checkpoint
# f32[512,512] from the argument W
# f32[256,512] from the argument x

The obvious next step is to ask the compiled program whether it needs less memory now, and on this backend the answer is no. memory_analysis().temp_size_in_bytes reports the same temporary arena with remat and without, at every depth tried: 2.0 MiB at two layers, 4.5 at four, 8.5 at eight, 18.0 at sixteen, identical in both columns.

Read what that number is before deciding it is wrong. The arena is what buffer assignment reserved for temporaries in the CPU executable. Dropping the residuals at the JAX level means the backward pass recomputes those values, and recomputed values need somewhere to live too, so the high-water mark the compiler reserved did not fall.

The lesson is about instruments, not about remat. The jaxprThe traced program: one equation per primitive in single-assignment form, every shape and dtype stated.taught in /l/jaxpr →-level list responded to the change and the compiled arena did not, and if you had only run the second one you would have concluded that jax.checkpoint does nothing. On an accelerator the number that decides whether a model fits is the allocator's peak, which comes from a device memory profile, and this CPU run cannot produce that figure. LAB·J3 covers the correctness half: both gradients agree.

Report the number the instrument actually produces, and name the instrument next to it.

jax.checkpoint with no policy saves nothing and recomputes everything. A policy names which values may be kept, and dots_saveable is the common one: keep the matmul outputs, recompute the cheap elementwise work around them.

On an eight-layer gelu stack that policy is the only variant that moved the arena, from 7.50 MiB down to 2.00, and it bought that with 0.35 percent more flops by the compiler's own count, 3.112G against 3.123G. On the four-layer tanh stack the identical policy moved neither number.

Two programs, one policy, opposite outcomes. Which values are worth keeping depends on what the backward pass actually needs and on what the compiler was going to do anyway, so a policy is a thing to measure per program rather than a setting to carry between them.

Set XLA_FLAGS=--xla_dump_to=DIR and compile, and the compiler writes the module out. Nothing has to execute: the script below calls .lower(...).compile() and prints a line saying so, and ten files land for that one module, including the HLO before and after optimization, the buffer assignment, and a memory usage report.

The report opens with the total the compiler reserved and then breaks it down by allocation. For a 256 by 256 tanh-matmul-sum it reserves 768.0 KiB: a 512.0 KiB temporary, the 256.0 KiB parameter, and a handful of four-byte allocations under that.

That is a memory answer available before you own the hardware, which is what makes the dump the first instrument to reach for when the question is structural. The XLAThe compiler: brilliant at fusing along dataflow edges, structurally unable to change your algorithm. That gap is why kernels exist.taught in /l/xla → path reads these same files pass by pass; here the point is only that they exist without a run.

$ cat s9.py
import jax
import jax.numpy as jnp

def step(p):
    return jnp.tanh(p @ p).sum()

jax.jit(step).lower(jnp.ones((256, 256))).compile()
print("compiled; nothing ran")

$ XLA_FLAGS=--xla_dump_to=dump2 python3 s9.py && ls dump2 | grep jit_step
compiled; nothing ran
module_0005.jit_step.before_optimizations.txt
module_0005.jit_step.cpu_after_optimizations-buffer-assignment.txt
module_0005.jit_step.cpu_after_optimizations-memory-usage-report.txt
module_0005.jit_step.cpu_after_optimizations.txt
module_0005.jit_step.ir-no-opt.part_00.ll
module_0005.jit_step.ir-no-opt.part_01.ll
module_0005.jit_step.ir-with-opt.part_00.ll
module_0005.jit_step.ir-with-opt.part_01.ll
module_0005.jit_step.obj-file.part_00.o
module_0005.jit_step.obj-file.part_01.o

$ head -9 dump2/module_0005.jit_step.cpu_after_optimizations-memory-usage-report.txt
Total bytes used: 786452 (768.0KiB)

Allocations sorted by size:

cumulative_size; total_size - cumulative_size; allocation
------------------------------------------------------------------------------
  512.0KiB( 67%);   256.0KiB; allocation 6: size 512.0KiB, preallocated-temp:
  768.0KiB(100%);        20B; allocation 0: size 256.0KiB, parameter 0, shape |f32[256,256]| at ShapeIndex {}:
  768.0KiB(100%);        16B; allocation 1: size 4B, output shape is |f32[]|, maybe-live-out:

A profiler trace is the other kind of answer, and it needs the program to actually run. jax.profiler.trace(dir) writes two files per capture, a gzipped Chrome-format JSON for Perfetto and an xplane protobuf for XProf, and TraceAnnotation puts your own label on a region so you can find it in either.

The capture is plain JSON underneath, so you can check what landed without opening a viewer. The events carry HLO op names with the module they came from, and the annotation is there as its own event. Two things in that output are worth reading carefully rather than quoting onward.

First, the annotated region measures 5.97 milliseconds while the ten steps it wraps take far longer, because the annotation covers the dispatch loop and the block sits outside it. Lesson one predicted exactly that shape. Second, the dot.3 events across all ten calls total 2.67 milliseconds, while the same ten steps take between 91.9 and 366.9 milliseconds of wall clock on this machine. Summing raw event durations from a CPU capture is not an accounting of the time; the work is spread across threadpool events, and aggregating them correctly is what XProf is for.

One more line to expect: the capture logs Can't import tensorflow.python.profiler.trace at error level when TensorFlow is absent, and writes both files regardless.

import collections
import glob
import gzip
import json

import jax
import jax.numpy as jnp

x = jnp.ones((1024, 1024))
step = jax.jit(lambda a: jnp.tanh(a @ a).sum())
step(x).block_until_ready()

with jax.profiler.trace("tr2"):
    with jax.profiler.TraceAnnotation("ten steps"):
        for _ in range(10):
            out = step(x)
    out.block_until_ready()

path = sorted(glob.glob("tr2/plugins/profile/*/*.trace.json.gz"))[-1]
events = json.load(gzip.open(path))["traceEvents"]
tracks = {e["pid"]: e["args"]["name"] for e in events if e.get("name") == "process_name"}
print(len(events), "events on tracks", sorted(tracks.values()))

total = collections.defaultdict(float)
count = collections.Counter()
for e in events:
    if e.get("ph") == "X" and not e["name"].startswith("$"):
        total[e["name"]] += e.get("dur", 0)
        count[e["name"]] += 1
for name, dur in sorted(total.items(), key=lambda kv: -kv[1])[:6]:
    print(f"{dur / 1000:8.2f} ms  x{count[name]:<5d} {name}")

# tr2/plugins/profile/2026_08_15_05_20_33/<host>.trace.json.gz
# tr2/plugins/profile/2026_08_15_05_20_33/<host>.xplane.pb
#
# 29497 events on tracks [/host:CPU]
#    77.25 ms  x10    ThunkExecutor::Execute (wait for completion)
#    11.70 ms  x20    PjitFunction(<lambda>)
#    10.99 ms  x10    call
#    10.94 ms  x10    tanh.4.clone
#     5.97 ms  x1     ten steps
#     3.52 ms  x10    call.1

Every instrument this arc reached for is in the table below, ordered by what it costs rather than by what it is called. The first four need no execution at all: a trace or a compile is enough, which is what lets you ask them about hardware you do not have in front of you.

The split worth carrying is the dump against the trace. A dump tells you what the compiler built, and a compile is the whole price. A trace tells you what the machine then did with that program, and it costs a real run on the real device. Ask the dump why a program has the shape it has; ask the trace where a run's time or memory actually went.