Chapter 2: Optimizing JAX Code for Performance
Course: "Advanced JAX: Performance, Optimization, and Scale", Chapter 2 · Original course link: https://apxml.com/courses/advanced-jax/chapter-2-optimizing-jax-code-performance/profiling-jax-code
Environment: Python 3.14 · JAX / jaxlib 0.12.0.1 · CPU only. GPU/TPU material is explained only; every runnable conclusion was verified on CPU.
Learning goals
- Generate traces on CPU with
jax.profilerand view them in Perfetto / TensorBoard; label traces and HLO withTraceAnnotationandnamed_scope. - Read jaxprs (
invars / constvars / eqns / outvars), walk a jaxpr programmatically and count primitives, and recognize how higher-order primitives such asjit,custom_jvp_callandcondnest. - Master the chain
jit(f).lower(...)→ StableHLO →.compile()→ optimized HLO, readcost_analysis()/memory_analysis(), and dump HLO withXLA_FLAGS=--xla_dump_to. - Understand row-major/column-major order, and when
jnp.transposereally moves data in XLA versus when it gets "absorbed". - Find and avoid repeated compilation, and use the persistent compilation cache so that a second process hits the cache directly.
- Count fusions in optimized HLO and explain why fusion saves memory bandwidth.
- Understand asynchronous dispatch, time code correctly, and know which operations synchronize implicitly.
- Optimize a real computation (a pairwise distance matrix) with the loop "implement → time → read jaxpr/HLO → refactor".
Sections
| Section | Content | Course page |
|---|---|---|
| 2.1 | Profiling JAX code on CPU, GPU and TPU | profiling-jax-code |
| 2.2 | Understanding the JAX computation graph (jaxpr) | understanding-jaxpr |
| 2.3 | The role of XLA compilation | role-of-xla-compilation |
| 2.4 | Memory layout and its effect on performance | memory-layout-performance |
| 2.5 | Avoiding recompilation | avoiding-recompilation |
| 2.6 | Operator fusion and op optimization | fusion-operator-optimization |
| 2.7 | Asynchronous dispatch | asynchronous-dispatch |
| 2.8 | Hands-on: optimizing a numerical computation | practice-optimizing-computation |
| — | Hands-on exercises · Quiz · Version notes · Summary · Cheat sheet |
How to read this
- This tutorial follows the 8 pages of Chapter 2 of the course. The explanations and code were rewritten from scratch, with extra material added. Outdated or error-prone patterns go in
❌cells, followed right away by the✅recommended version. - All conclusions are based on what JAX 0.12.0.1 actually does. Points that need special attention are marked with
> ⚠️ **Note**, and the "Version notes" at the end of the chapter collects them.
How this maps to the Basics part
- Basics 02 JIT compilation:
jitbasics, the full table of recompilation triggers,jax_log_compiles,_cache_size(), an intro to AOT (lower/compile),static_argnums,donate_argnums. This chapter doesn't repeat that material; it builds on it (persistent cache,explain_cache_misses, weak types, function object identity, and so on). - Basics 01 Introduction to JAX: an introduction to asynchronous dispatch and
block_until_ready. - Basics 07 Practical topics: a performance troubleshooting workflow in practice.
⚠️ Known issue in this environment (be sure to read): the CPU backend of jaxlib 0.12.0.1 has a bug that silently returns wrong results: if a jitted function captures a fairly large constant matrix through a closure, does a matmul with it and then reduces, the result can be wrong. So in this tutorial weights and data are always passed in as function arguments, and key numbers are cross-checked against NumPy float64. Section 2.3 dissects this bug as a real-world case of "why you should inspect the compiled output".
0. Environment and shared helpers¶
The cell below defines small helpers used throughout the chapter:
bench: warms up first, then times several runs and averages them, callingblock_until_readyevery time;clean_hlo: strips the source-location tables andmetadatafrom HLO text so it's easier to read;count_fusions: counts thefusion(...)instructions in optimized HLO;CompileCounter: used withjax.log_compiles()to count the XLA compilations that actually happen, so we can writeasserts for claims about recompilation;run_py: runs a script in a subprocess.XLA_FLAGSis read only once at process startup, and verifying the persistent cache means checking "whether a second process can hit it"; both of these can only be demonstrated in a subprocess.
import os, re, sys, time, gzip, glob, json, logging, tempfile, textwrap, subprocess, functools
from collections import Counter
import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
import jax.extend as je
import jaxlib
print("JAX", jax.__version__, "| jaxlib", jaxlib.__version__, "| backend:", jax.default_backend(), "|", jax.devices())
def bench(f, *args, n=20):
"""Warm up first (including compilation), then return the mean time of n runs (seconds); block_until_ready every time."""
jax.block_until_ready(f(*args))
t0 = time.perf_counter()
for _ in range(n):
jax.block_until_ready(f(*args))
return (time.perf_counter() - t0) / n
_TABLE_HDR = re.compile(r'^(FileNames|FunctionNames|FileLocations|StackFrames)$')
def clean_hlo(text):
"""Strip source-location tables, metadata and blank lines from HLO, leaving just the computation."""
out = []
for line in text.splitlines():
if _TABLE_HDR.match(line) or re.match(r'^\d+ ', line) or not line.strip():
continue
out.append(re.sub(r', metadata=\{[^}]*\}', '', line))
return "\n".join(out)
def count_fusions(hlo_text):
"""Count the fusion instructions (= fused "kernels") in optimized HLO."""
return len(re.findall(r' fusion\(', hlo_text))
def compiled_hlo(f, *args, **jit_kw):
"""Return the compiled HLO text of jit(f) for args."""
return jax.jit(f, **jit_kw).lower(*args).compile().as_text()
class CompileCounter(logging.Handler):
"""with CompileCounter("f") as cc: ... afterwards cc.count is the number of times jit(f) actually compiled."""
def __init__(self, name=None):
super().__init__()
self.name, self.records = name, []
def emit(self, record):
msg = record.getMessage()
if msg.startswith("Compiling jit("):
fn = msg[len("Compiling jit("):].split(")")[0]
if self.name is None or fn == self.name:
self.records.append(msg)
def __enter__(self):
self._logger = logging.getLogger("jax")
self._logger.addHandler(self)
# Temporarily raise the level of jax's own StreamHandler so log_compiles doesn't flood the output
self._saved = [(h, h.level) for h in self._logger.handlers if h is not self]
for h, _ in self._saved:
h.setLevel(logging.ERROR)
self._ctx = jax.log_compiles(True)
self._ctx.__enter__()
return self
def __exit__(self, *exc):
self._ctx.__exit__(*exc)
for h, lvl in self._saved:
h.setLevel(lvl)
self._logger.removeHandler(self)
@property
def count(self):
return len(self.records)
SCRATCH = tempfile.mkdtemp(prefix="adv02_")
def run_py(src, env_extra=None, args=()):
"""Write src to a script, run it in a fresh Python subprocess, and return the stdout+stderr text."""
path = os.path.join(SCRATCH, f"sub_{abs(hash(src)) % 10**8}.py")
with open(path, "w") as fh:
fh.write(textwrap.dedent(src))
env = dict(os.environ)
env.update(env_extra or {})
p = subprocess.run([sys.executable, path, *args], capture_output=True, text=True, env=env, timeout=300)
if p.returncode != 0:
raise RuntimeError(p.stdout + p.stderr)
return p.stdout + p.stderr
print("temp directory:", SCRATCH)
2.1 Profiling JAX code on CPU, GPU and TPU¶
jax.jit is only the starting point for speed; to get the best performance you have to profile. Code that's fast on CPU can behave completely differently on GPU/TPU, for three reasons: the hardware architectures differ, the memory bandwidth limits differ, and XLA applies different optimizations for different hardware.
Profiling mainly answers five questions:
- Which operations take the most time?
- Is the program compute-bound or memory-bound?
- Is the accelerator fully used? Are there idle stretches?
- Are there unexpected host ↔ device data transfers slowing everything down?
- Is there too much recompilation?
Collecting a trace with jax.profiler¶
The core APIs are jax.profiler.start_trace(log_dir) and jax.profiler.stop_trace(). The course example goes like this:
- Define a
@jax.jitfunctioncomplex_computation(x, y): firstjnp.dot, thenjnp.sin, and finally returnjnp.mean(z * x + y); - Generate two 2000×2000 random matrices and put them on the device with
jax.device_put; - Call it once and
block_until_ready()as a warm-up, so compilation stays out of the trace; start_trace(log_dir)sets where the trace is saved;- Run it several times in a loop (the course uses 5) so the trace is representative;
block_until_ready()every time. JAX dispatches asynchronously (see 2.7); without blocking, Python might stop the trace before the computation finishes;- End with
stop_trace(), which writes to disk.
The old style generates data with jax.random.PRNGKey(0) and has x and y share the same key. Let's look at that first, then the recommended version (for an intro to the two kinds of keys, see "Random numbers" in Basics 01 Introduction to JAX; for more depth, see Advanced 10 Random numbers and sharding).
# ❌ Old style (outdated): jax.random.PRNGKey returns a key as a "raw uint32 array"; 0.12 recommends typed keys
old_key = jax.random.PRNGKey(0)
new_key = jax.random.key(0)
print("PRNGKey(0):", old_key, "| dtype =", old_key.dtype, "| shape =", old_key.shape)
print("key(0) :", new_key.dtype, "| shape =", new_key.shape)
try:
# A typed key can't be used in arithmetic like a plain integer array, which prevents mistaking the key for data
_ = new_key + 1
print("a typed key can be added to?")
except Exception as e:
print("arithmetic on a typed key raises:", type(e).__name__)
assert old_key.dtype == jnp.uint32 and old_key.shape == (2,)
assert jnp.issubdtype(new_key.dtype, jax.dtypes.prng_key) and new_key.shape == ()
# ✅ Recommended: jax.random.key + split, with different subkeys for x and y
@jax.jit
def complex_computation(x, y):
z = jnp.dot(x, y)
z = jnp.sin(z)
return jnp.mean(z * x + y)
kx, ky = jax.random.split(jax.random.key(0))
x = jax.device_put(jax.random.normal(kx, (2000, 2000)))
y = jax.device_put(jax.random.normal(ky, (2000, 2000)))
# Warm-up: trigger compilation and wait for it to really finish
complex_computation(x, y).block_until_ready()
log_dir = tempfile.mkdtemp(prefix="jax_profiling_") # the original example uses /tmp/jax_profiling
jax.profiler.start_trace(log_dir)
for _ in range(5):
result = complex_computation(x, y)
result.block_until_ready() # important: block so the computation finishes before profiling stops
jax.profiler.stop_trace()
print("trace saved to:", log_dir)
files = sorted(p for p in glob.glob(log_dir + "/**", recursive=True) if os.path.isfile(p))
for p in files:
print(f" {os.path.relpath(p, log_dir):70s} {os.path.getsize(p):>8d} bytes")
assert any(p.endswith(".xplane.pb") for p in files)
assert any(p.endswith(".trace.json.gz") for p in files)
⚠️ Note: The trace file is not a single
profile.pb.start_trace/tracecreates two files underlog_dir/plugins/profile/<timestamp>/:<hostname>.xplane.pb(XSpace protobuf, the format TensorBoard/XProf reads) and<hostname>.trace.json.gz(Chrome trace format). If you passcreate_perfetto_trace=True, you also get aperfetto_trace.json.gz. Theasserts above already verify the first two files.
How to view a trace¶
| Tool | How | Notes |
|---|---|---|
| Perfetto (recommended, nothing to install) | open https://ui.perfetto.dev → Open trace file → pick perfetto_trace.json.gz (or *.trace.json.gz) | parsed locally in the browser; the file isn't uploaded |
create_perfetto_link=True | jax.profiler.trace(dir, create_perfetto_link=True) | prints a link at the end and blocks waiting for a browser to open it; fine for scripts, not for batch-running in a notebook |
| TensorBoard / XProf | pip install tensorboard xprof (formerly tensorboard-plugin-profile), then tensorboard --logdir <log_dir> and open the Profile tab | installing only tensorboard doesn't give you the Profile tab. Not installed in this environment |
The course covers three views in TensorBoard Profile:
- Trace Viewer: the most informative one.
- Host threads: Python calls, JAX dispatch overhead, and any NumPy operations.
- Device streams: e.g.
/GPU:0/stream:all,/TPU:0/stream:all.HtoDis a host-to-device copy andDtoHa device-to-host copy. - What to look for: long kernels on device streams (compute bottlenecks), large gaps (device idle), and long copies (data transfer bottlenecks). In the course's illustrative timeline, the GPU compute stream has a gap between 2.5 and 4 ms, meaning the device is idle or waiting for data.
- Ops View: totals per operation type (e.g.
dot_general,sin,reduce_mean): total time, average time and call count. - Memory Viewer: shows how memory allocation changes over time. A very high peak or frequent allocation/freeing can both be problems. The availability and level of detail of this view vary by backend.
Below we try a different style: the context manager jax.profiler.trace, together with TraceAnnotation (names a time span on the host side) and jax.named_scope (adds a name prefix to operations in the computation graph). Then we parse the trace file right in the notebook to confirm the labels really got written.
@jax.jit
def annotated_step(x, y):
with jax.named_scope("my_matmul"): # applies to op names in the jaxpr/HLO
z = x @ y
with jax.named_scope("my_epilogue"):
out = jnp.mean(jnp.sin(z))
return out
a = jax.random.normal(kx, (512, 512)); b = jax.random.normal(ky, (512, 512))
annotated_step(a, b).block_until_ready() # warm-up
log_dir2 = tempfile.mkdtemp(prefix="jax_trace_")
with jax.profiler.trace(log_dir2, create_perfetto_trace=True):
for i in range(3):
with jax.profiler.TraceAnnotation(f"step_{i}"): # applies to the host-side timeline
annotated_step(a, b).block_until_ready()
files2 = sorted(os.path.relpath(p, log_dir2) for p in glob.glob(log_dir2 + "/**", recursive=True) if os.path.isfile(p))
print("\n".join(files2))
assert any(f.endswith("perfetto_trace.json.gz") for f in files2)
# trace.json.gz is standard Chrome trace format: {"traceEvents": [...]}
tj = glob.glob(log_dir2 + "/**/*.trace.json.gz", recursive=True)[0]
events = json.load(gzip.open(tj))["traceEvents"]
names = Counter(e["name"] for e in events if e.get("ph") == "X")
print("\nTop 12 event names by count:")
for k, v in names.most_common(12):
print(f" {v:4d} {k}")
assert all(names[f"step_{i}"] == 1 for i in range(3)) # TraceAnnotation made it into the trace
assert names["PjitFunction(annotated_step)"] >= 3 # host-side event for each call of the jitted function
# xplane.pb can also be parsed directly with jax.profiler.ProfileData (no TensorBoard needed)
xp = glob.glob(log_dir2 + "/**/*.xplane.pb", recursive=True)[0]
pdata = jax.profiler.ProfileData.from_file(xp)
for plane in pdata.planes:
lines = list(plane.lines)
print(f"plane {plane.name!r}: {len(lines)} lines")
for ln in lines:
if "XLAExecute" in ln.name or ln.name == "":
evs = [e.name for e in ln.events]
kern = sorted({n for n in evs if "fusion" in n or n.startswith("wrapped")})
print(f" line {ln.name!r:28s} {len(evs):3d} events; kernel names among them: {kern}")
# named_scope doesn't change the computation, but it shows up in HLO op_name metadata (handy when debugging/matching against a trace)
lowered = annotated_step.lower(a, b)
hlo_ops = re.findall(r'op_name="([^"]*my_[^"]*)"', lowered.compile().as_text())
print("\nHLO op_names with a named_scope prefix:", sorted(set(hlo_ops)))
assert any("my_matmul" in s for s in hlo_ops) and any("my_epilogue" in s for s in hlo_ops)
assert "my_matmul" not in lowered.as_text() # the default StableHLO text has no location info
assert "my_matmul" in lowered.as_text(debug_info=True) # only shown with debug_info=True
Notes for different devices¶
- CPU:
jax.profilercaptures the parts that run on CPU after jit compilation. Pure Python parts (data loading, preprocessing loops outside jit, script logic) are still best profiled withcProfile. Alternating between jitted functions and lots of Python logic adds overhead. - GPU: the Trace Viewer separates compute streams from copy streams (PCIe transfers). Focus on how long
HtoD/DtoHtake, and ask yourself two questions: am I moving data I don't need to? Can the data stay on the GPU longer? For JAX-level optimization, TensorBoard is usually enough; for per-kernel detail, use NVIDIA Nsight Systems / Nsight Compute. - TPU: focus on the utilization percentage of the MXU (matrix multiply unit). To make dimensions fit the MXU's requirements (e.g. divisible by 128), XLA may pad, and that sometimes shows up as overhead. Low utilization can come from the input pipeline, the host CPU, or the computation itself being a poor fit for TPU. On a TPU VM you can use Google Cloud Profiler for system-level information.
Common bottlenecks¶
- Host–device data transfer: long
HtoD/DtoHbars on the copy stream. The fix is to move data to the device once and keep it there as much as possible, and to batch. - Small kernels / launch latency: lots of very short bars on the device stream. XLA fuses as much as it can, but the computation may still be chopped too finely.
- Device idle: obvious gaps on the compute stream timeline. Maybe the host can't feed data fast enough, transfers are blocking, or there are synchronization points.
- Compute-bound operations: a few very long bars fill the compute stream. You may need a different algorithm, a more efficient primitive, or mixed precision.
Device memory profile¶
jax.profiler.save_device_memory_profile(path) aggregates the device buffers that are currently alive by allocation site (Python call stack) and writes them in pprof format (gzip-compressed protobuf). First let's confirm it works on CPU:
mem_path = os.path.join(SCRATCH, "memory.prof")
blob_before = jax.profiler.device_memory_profile() # returns bytes, doesn't write to disk
big = jnp.ones((4000, 4000)); big.block_until_ready() # add a live buffer of about 64 MB
jax.profiler.save_device_memory_profile(mem_path) # the version that writes to disk
blob_after = open(mem_path, "rb").read()
print("works on CPU; file size:", os.path.getsize(mem_path), "bytes; gzip magic:", blob_after[:2])
print("decompressed protobuf size: before", len(gzip.decompress(blob_before)), "→ after", len(gzip.decompress(blob_after)))
assert blob_after[:2] == b"\x1f\x8b" # pprof in gzip format
del big
print("\nTo view: pprof -http=:8080", mem_path, " (or go tool pprof ...)")
To view it: use the Go version of pprof (go install github.com/google/pprof@latest) and run pprof -http=:8080 memory.prof, then check Flame Graph / Top to see which line of Python holds the most device memory. It's a snapshot at one moment, not a time series, which makes it good for answering "who is holding the device memory".
🔍 Extra¶
⚪ Optional: read this when you want to match kernel names in a trace or label training steps.
- A CPU trace has only one plane,
/host:CPU. XLA's execution threads are calledtf_XLAExecute/…, and the kernel names are the names of fusion instructions in the optimized HLO (theynn_fusion,wrapped_multiply, etc. seen above), so names in the trace can be looked up directly in the HLO (see 2.3). TraceAnnotationlabels host-side time spans.named_scopelabels op names in the computation graph, goes into HLO'sop_name, and is visible in GPU/TPU traces too. They serve different purposes and can't replace each other.- For training loops you can use
jax.profiler.StepTraceAnnotation("train", step_num=i), and TensorBoard will split the statistics by step;@jax.profiler.annotate_functionis the decorator version ofTraceAnnotation. - Always warm up before tracing, or the first step's tracing and compilation will take up most of the timeline. Keep traces short, a few steps is enough; otherwise the file gets huge and the browser can't open it.
2.2 Understanding the JAX computation graph (jaxpr)¶
jit compiles functions for accelerators. There's an intermediate step: the Python code first becomes a jaxpr (JAX program representation), an intermediate representation of the computation graph, which is then handed to XLA. Understanding jaxprs helps you see what JAX actually traced, and helps you find the root cause of inefficiency or recompilation.
Where a jaxpr comes from: the first time a jitted function is called with a given input type/shape, JAX traces it. It replaces the real inputs with tracer objects and records every primitive operation performed on those tracers; the result is the jaxpr. This step doesn't actually run the computation.
Four properties of a jaxpr:
- Functional: it's a pure function with no side effects.
- Statically typed: every variable has a definite shape and dtype.
- Explicit operations: a jaxpr is a sequence of equations, each applying a primitive (such as
add,sin,dot_general,reduce_sum) to variables or constants. Pythonif/forare unrolled at tracing time, unless you use JAX control-flow primitives likelax.condandlax.scan. - Intermediate representation: it sits between Python source and low-level XLA HLO.
Generating and inspecting a jaxpr¶
The main tool is jax.make_jaxpr(f)(*example_inputs). Here's the course example:
def example_function(x, y):
a = jnp.sin(x)
b = jnp.cos(y)
return a + b * 2.0
kx_, ky_ = jax.random.split(jax.random.key(0)) # using the recommended jax.random.key (the old style is PRNGKey)
x_example = jax.random.normal(kx_, (10,))
y_example = jax.random.normal(ky_, (10,))
closed = jax.make_jaxpr(example_function)(x_example, y_example)
print(closed.jaxpr)
Item by item:
{ lambda ; ... }: defines a lambda. Left of the semicolon are theconstvars(closure constants), right of it are theinvars(input variables);a:f32[10] b:f32[10]: input variables with their dtype and shape;let ... in ...: each line afterletis an equation (eqn), in the formoutput_var:type = primitive[params] inputs...;in (f,): the output variables,outvars.
0.12 prints literals with their type, e.g. mul d 2.0:f32[]. Some materials show mul d 2.0; that's only a display difference.
The object returned by make_jaxpr carries both the jaxpr itself and the actual values of the constants, consts. In older versions it was a ClosedJaxpr wrapping a Jaxpr, and you used .jaxpr to get at the inner one; JAX 0.12 merged the two into a single Jaxpr class. jax.extend.core.ClosedJaxpr is just an alias of Jaxpr, .jaxpr is kept as a compatibility attribute (returning the object itself), and .consts and .eqns are directly accessible. So closed.jaxpr.eqns in old code still works (in the output below, both type names are Jaxpr). Let's take the four fields apart in code:
j = closed.jaxpr
print(type(closed).__name__, "→ .jaxpr is", type(j).__name__)
print("invars :", [(str(v), v.aval.str_short()) for v in j.invars])
print("constvars:", j.constvars, "| consts:", closed.consts)
print("outvars :", [(str(v), v.aval.str_short()) for v in j.outvars])
for eqn in j.eqns:
print(f" eqn: {eqn.primitive.name:4s} in={[str(v) for v in eqn.invars]} out={[str(v) for v in eqn.outvars]} params={eqn.params}")
# 0.12: ClosedJaxpr has been merged into Jaxpr; .jaxpr is a compatibility attribute that returns the object itself
assert type(closed) is je.core.ClosedJaxpr is type(j) and closed.jaxpr is closed
assert [e.primitive.name for e in j.eqns] == ["sin", "cos", "mul", "add"]
assert len(j.invars) == 2 and len(j.outvars) == 1 and j.constvars == []
# The literal 2.0 isn't a variable but a Literal
assert type(j.eqns[2].invars[1]).__name__ == "Literal"
Where constvars come from: JAX arrays captured by the function's closure aren't arguments but constants, and they show up left of the semicolon:
W_small = jnp.arange(9.0).reshape(3, 3) # a small matrix, just to show the jaxpr structure
def uses_closure(v):
return jnp.dot(W_small, v) + jnp.pi
cj = jax.make_jaxpr(uses_closure)(jnp.ones(3))
print(cj)
print("constvars:", [v.aval.str_short() for v in cj.jaxpr.constvars], "| consts shapes:", [c.shape for c in cj.consts])
assert len(cj.jaxpr.constvars) == 1 and cj.consts[0].shape == (3, 3)
# The Python scalar jnp.pi becomes a literal directly
assert type(cj.jaxpr.eqns[1].invars[1]).__name__ == "Literal"
Walking a jaxpr: counting primitives and seeing how higher-order primitives nest¶
eqns is an ordinary Python list, so you can iterate over it directly. Higher-order primitives (jit, cond, while, scan, custom_jvp_call, custom_vjp_call, …) hide sub-jaxprs in their params, so you need to recurse into them. jax.extend.core.jaxprs_in_params(eqn.params) finds these sub-jaxprs for you.
def count_primitives(jaxpr, counter=None, depth=0, show=False):
"""Recursively count primitive occurrences in a jaxpr (including all sub-jaxprs)."""
counter = Counter() if counter is None else counter
for eqn in jaxpr.eqns:
counter[eqn.primitive.name] += 1
subs = list(je.core.jaxprs_in_params(eqn.params))
if show and subs:
print(" " * depth + f"└ {eqn.primitive.name} (contains {len(subs)} sub-jaxprs)")
for sub in subs:
count_primitives(sub, counter, depth + 1, show)
return counter
@jax.jit
def inner(x):
return jax.nn.relu(x) * 2.0
def outer(x):
return inner(x).sum() + jnp.logaddexp(x, 1.0).sum()
co = jax.make_jaxpr(outer)(jnp.ones(3))
print("top-level primitives:", [e.primitive.name for e in co.jaxpr.eqns])
print("nesting structure:")
cnt = count_primitives(co.jaxpr, show=True)
print("recursive counts:", dict(cnt))
top = co.jaxpr.eqns[0]
assert top.primitive.name == "jit" and top.params["name"] == "inner"
# In 0.12 the primitive for nested jit is named "jit" (older versions called it "pjit")
assert "pjit" not in cnt
# relu itself is a custom_jvp function that wraps another jit inside; logaddexp is custom_jvp too
assert cnt["custom_jvp_call"] == 2 and cnt["jit"] == 2 and cnt["max"] >= 1
Why jaxprs matter for optimization¶
- Tracing visibility: you can spot whether JAX traced more computation than you expected. Subtle differences between NumPy and
jax.numpycan produce different primitives. - Understanding transformations: look at the jaxpr after applying
jax.gradorjax.vmapto a function, and you can see how the graph was rewritten. - Diagnosing recompilation: compare the jaxprs produced by different calls; if the
eqnslists differ noticeably, the graph depends on something "dynamic".
⚠️ Note: You'll often hear: "the jaxpr structure changes between calls, causing
jax.jitto recompile, usually because the sequence of operations depends on argument values (Python control flow based on tensor values)". In fact, a Pythonifon the value of a traced argument insidejitcan't be traced at all: it raisesTracerBoolConversionErrorright away and never gets as far as "recompilation". What triggers recompilation is a change in the cache key: input shapes, dtypes, weak types, the values ofstatic_argnumsarguments, the function object itself, and so on (see 2.5). Only when an argument is marked static do you get "different value → different jaxpr → recompilation". The next two cells verify this.
- Connection to XLA: XLA's optimizations (such as fusion) are affected by the structure of the input jaxpr. A clear, predictable jaxpr is easier to optimize.
# ❌ Common misconception: "value-dependent Python control flow changes the jaxpr and so triggers recompilation"
@jax.jit
def value_branch(x, flag):
if flag > 0.5: # takes bool of a tracer
return jnp.sin(x)
return jnp.cos(x)
try:
value_branch(jnp.ones(3), 0.7)
caught = None
except Exception as e:
caught = type(e).__name__
print("measured:", caught, "|", str(e).splitlines()[0])
assert caught == "TracerBoolConversionError"
# ✅ Recommended: when a "value" really should decide the graph structure, mark it static; then different values → different jaxprs → each compiled once
def value_branch_static(x, flag):
return jnp.sin(x) if flag > 0.5 else jnp.cos(x)
jp_hi = jax.make_jaxpr(value_branch_static, static_argnums=1)(jnp.ones(3), 0.7)
jp_lo = jax.make_jaxpr(value_branch_static, static_argnums=1)(jnp.ones(3), 0.3)
print(jp_hi, jp_lo, sep="\n")
assert [e.primitive.name for e in jp_hi.jaxpr.eqns] == ["sin"]
assert [e.primitive.name for e in jp_lo.jaxpr.eqns] == ["cos"]
jit_static = jax.jit(value_branch_static, static_argnums=1)
with CompileCounter("value_branch_static") as cc:
for flag in [0.7, 0.7, 0.3, 0.3]:
jit_static(jnp.ones(3), flag)
print("actual compilations over 4 calls:", cc.count)
assert cc.count == 2
An example with control flow: lax.cond¶
In the course example conditional_function(use_sin, x), pred = use_sin > 0.5, the true branch is jnp.sin, the false branch is jnp.cos, and it returns lax.cond(pred, true_fun, false_fun, x). The predicate of lax.cond must be a scalar. Tracing it with 0.7 and with 0.3, you can see:
- Both traces give exactly the same jaxpr structure. The jaxpr doesn't contain only the selected branch; it contains one
condprimitive with both branches packaged inside; - As long as both branches have matching input/output types and shapes, the predicate's value doesn't affect the jaxpr structure. This matters for jit, because the same compiled program has to handle either branch.
def conditional_function(use_sin, x):
pred = use_sin > 0.5 # scalar boolean
def true_fun(v):
return jnp.sin(v)
def false_fun(v):
return jnp.cos(v)
return lax.cond(pred, true_fun, false_fun, x)
xs3 = jnp.arange(3.0)
jc_07 = jax.make_jaxpr(conditional_function)(0.7, xs3)
jc_03 = jax.make_jaxpr(conditional_function)(0.3, xs3)
print(jc_07)
assert str(jc_07) == str(jc_03) # structure doesn't depend on the predicate's value
cond_eqn = [e for e in jc_07.jaxpr.eqns if e.primitive.name == "cond"][0]
branch_prims = [[e.primitive.name for e in br.jaxpr.eqns] for br in cond_eqn.params["branches"]]
print("branches order:", branch_prims)
prev = jc_07.jaxpr.eqns[1]
print("the predicate is converted first:", prev.primitive.name, prev.params["new_dtype"])
assert branch_prims == [["cos"], ["sin"]] # branch 0 is false_fun
assert prev.primitive.name == "convert_element_type" and prev.params["new_dtype"] == jnp.int32
⚠️ Note: In a jaxpr,
conddoesn't take the formcond c true_branch false_branch b, i.e. "true branch first, then false branch, using the boolean predicate directly". The boolean predicate first goes throughconvert_element_typeto become an int32 index, andcond's parameter isbranches=(false_fun, true_fun). So index 0 is the false branch and index 1 the true branch, the same rule aslax.switch. Theassertabove already verifies this.
Limitations¶
A jaxpr can't show XLA's final optimizations (such as operator fusion), nor the low-level code on the target accelerator. It's an important tool for understanding and debugging JAX internals, but it's just one intermediate step in the compilation pipeline. The next section shows how to look further down.
🔍 Extra¶
⚪ Optional: read this when you want to look at transformed jaxprs or write your own jaxpr interpreter.
- See what a transformation did:
jax.make_jaxpr(jax.grad(f))(x)shows the forward and backward computation directly. In the example below, the forwardsinis kept for reuse in the backward pass (thecoscomes from the derivative ofsin). - A jaxpr can be executed as a function:
jax.extend.core.jaxpr_as_fun(closed_jaxpr)(*args)interprets the jaxpr equation by equation, which is handy for writing custom interpreters or for verification. - A common misconception: a large intermediate like
f32[1000,1500,64]in a jaxpr doesn't mean that much memory will necessarily be allocated at runtime. After XLA fusion it may never be materialized at all; Section 2.8 measures this. To judge memory use, look atmemory_analysis()after compilation.
def f_grad_demo(x):
return jnp.sum(jnp.sin(x) ** 2)
gj = jax.make_jaxpr(jax.grad(f_grad_demo))(jnp.ones(4))
print(gj)
prims = count_primitives(gj.jaxpr)
assert prims["sin"] == 1 and prims["cos"] == 1 # the backward pass uses d sin = cos
# jaxpr_as_fun: call the ClosedJaxpr like a regular function; the result should match the original function
fn_from_jaxpr = je.core.jaxpr_as_fun(gj)
np.testing.assert_allclose(fn_from_jaxpr(jnp.arange(4.0))[0], jax.grad(f_grad_demo)(jnp.arange(4.0)), rtol=1e-6)
print("jaxpr_as_fun result matches jax.grad")
2.3 The role of XLA compilation¶
jax.jit is the entry point, but what actually optimizes for the hardware and generates code is XLA (Accelerated Linear Algebra). XLA is a domain-specific compiler developed by Google that specializes in numerical computation dominated by linear algebra, and it supports CPU, GPU and TPU. You can think of jax.jit as the mechanism that "hands the computation graph to XLA", and XLA optimizes it and generates executable code.
XLA works on a high-level intermediate representation called HLO (High Level Operations). It first runs a series of hardware-independent and hardware-specific optimization passes over the HLO graph, then compiles it into machine code for the target device.
⚠️ Note: A common description is "JAX converts the traced jaxpr into HLO", and some people expand HLO as "High-Level Optimizer IR". To be precise, the jaxpr is first lowered to StableHLO (an MLIR dialect; what
jit(f).lower(...).as_text()prints isstablehlo.*ops). XLA then takes over, converts it into HLO, and optimizes on HLO (.compile().as_text()prints the optimized HLO, which starts withHloModule). In current XLA documentation, HLO stands for High Level Operations. The cell below verifies the difference between these two levels of text withassert.
The full pipeline is:
$$ \text{Python function} \xrightarrow{\text{trace}} \text{jaxpr} \xrightarrow{\text{lower}} \text{StableHLO} \xrightarrow{\text{XLA optimization passes}} \text{optimized HLO} \xrightarrow{\text{backend}} \text{machine code} $$
The corresponding AOT API is jax.jit(f).trace(...) → .lower(...) → .compile(). Basics 02 JIT compilation already showed how to split these stages; here the focus is on reading the outputs.
def simple_computation(x, a, b):
y = jnp.sin(x)
z = a * y
w = z + b
return w
xs = jnp.ones((1000, 1000))
lowered = jax.jit(simple_computation).lower(xs, 2.0, 1.0)
stablehlo_txt = lowered.as_text()
print("==== StableHLO (after lowering, before XLA optimization) ====")
print(stablehlo_txt)
compiled = lowered.compile()
opt_hlo = compiled.as_text()
print("==== Optimized HLO (after compile) ====")
print(clean_hlo(opt_hlo))
assert "stablehlo.sine" in stablehlo_txt and "HloModule" not in stablehlo_txt
assert opt_hlo.startswith("HloModule")
# The three elementwise ops (sin / mul / add) are fused into one fusion instruction
assert count_fusions(opt_hlo) == 1
print("fusions after optimization:", count_fusions(opt_hlo))
This is exactly the operator fusion the course talks about: without fusion you'd need 3 kernels, and the intermediates y and z would both be written back to memory. After fusion there's just one kernel that does sin, the multiply and the add in one go; intermediate values stay in registers and only the final w is written back.
cost_analysis() and memory_analysis()¶
The compiled object can report XLA's estimates of compute and memory use:
ca = compiled.cost_analysis()
ma = compiled.memory_analysis()
print("cost_analysis type:", type(ca).__name__)
for k in ["flops", "transcendentals", "bytes accessed"]:
print(f" {k:16s} = {ca[k]:,.0f}")
print("memory_analysis:")
for k in ["argument_size_in_bytes", "output_size_in_bytes", "temp_size_in_bytes", "alias_size_in_bytes"]:
print(f" {k:24s} = {getattr(ma, k):,}")
n_el = 1000 * 1000
assert ca["flops"] == 2 * n_el # one multiply + one add
assert ca["transcendentals"] == n_el # one sin
# After fusion it only reads x and writes w (plus two scalars): two passes of 4 MB
assert ca["bytes accessed"] == 2 * 4 * n_el + 8
assert ma.output_size_in_bytes == 4 * n_el and ma.temp_size_in_bytes == 0
- In 0.12,
cost_analysis()returns a dict (older versions returned a[dict]list). Keys includeflops,transcendentals,bytes accessed, and per-operand breakdowns such asbytes accessed0{}. It's XLA's static estimate, not a measurement. memory_analysis()returnsCompiledMemoryStats, wheretemp_size_in_bytesis the size of intermediate buffers. It's the most direct evidence of whether "a large intermediate really got materialized" (used in 2.6 and 2.8).
Five kinds of XLA optimization¶
- Operator fusion: the course considers this the most impactful optimization; see the example above and 2.6.
- Constant folding: parts that depend only on compile-time constants are computed at compile time.
- Algebraic simplification: e.g. simplifying
x * 1.0tox. - Layout optimization: XLA picks how arrays are arranged in memory based on the computation and the hardware (see 2.4).
- Target-specific code generation: after the hardware-independent optimizations, it generates low-level code for CPU/GPU/TPU. CPU and GPU go through LLVM; TPU uses a dedicated compiler.
First let's verify constant folding and algebraic simplification:
# Constant folding: jnp.sin(2.0) and arange(4).sum() depend only on constants → already computed in the optimized HLO
fold = lambda v: v * jnp.sin(2.0) + jnp.arange(4.0).sum()
v4 = jnp.ones(4)
low_txt = jax.jit(fold).lower(v4).as_text()
opt_txt = compiled_hlo(fold, v4)
print("sine / reduce in StableHLO:", "stablehlo.sine" in low_txt, "stablehlo.reduce" in low_txt)
print("sine / reduce in optimized HLO:", " sine(" in opt_txt, " reduce(" in opt_txt)
assert "stablehlo.sine" in low_txt and " sine(" not in opt_txt and " reduce(" not in opt_txt
# Algebraic simplification ①: x * 1.0 + 0.0 → just copy x
simp = clean_hlo(compiled_hlo(lambda v: v * 1.0 + 0.0, xs))
print("\nx*1.0+0.0 after optimization:\n", simp.split("ENTRY")[-1])
assert " multiply(" not in simp and " add(" not in simp
# ❌ Common misconception: thinking XLA simplifies (x + y) - x to y
alg = clean_hlo(compiled_hlo(lambda p, q: (p + q) - p, xs, xs))
print("\n".join(l for l in alg.splitlines() if " add(" in l or " subtract(" in l))
try:
assert " subtract(" not in alg, "the HLO still contains subtract"
print("it was simplified")
except AssertionError as e:
print("measured: not simplified --", e)
# ✅ Correct: in floating point (x+y)-x ≠ y (precision is lost), so XLA doesn't do this value-changing rewrite by default
p_, q_ = jnp.float32(1e8), jnp.float32(1.0)
exact_y = q_
via_xla = jax.jit(lambda p, q: (p + q) - p)(p_, q_)
print("(1e8 + 1) - 1e8 =", via_xla, "(mathematically it should be 1)")
assert " subtract(" in alg and float(via_xla) != float(exact_y)
⚠️ Note:
(x + y) - x → yis often used as an example of algebraic simplification, with at most a remark that "it doesn't always hold exactly because of floating point". In fact XLA does not simplify this expression (JAX 0.12, CPU); the optimized HLO still containsaddandsubtract. The reason is that under IEEE floating point this rewrite changes the result. For example,(1e8+1)-1e8is 0 in float32, whileyis 1. By default, XLA's algebraic simplifier doesn't do rewrites that change IEEE results. The simplifications we saw above,x*1.0 → xandx+0.0 → x, are all value-preserving.
Dumping HLO at every stage with XLA_FLAGS=--xla_dump_to¶
XLA_FLAGS is read only once, when the process starts and XLA initializes, so we demonstrate it in a subprocess. In the dump directory, each compiled module produces a set of files. The ones you'll look at most are:
*.before_optimizations.txt: the HLO as XLA receives it, before fusion;*.cpu_after_optimizations.txt: the optimized HLO, the same as.compile().as_text();*-buffer-assignment.txt/*-memory-usage-report.txt: buffer assignment and memory report.
dump_dir = os.path.join(SCRATCH, "xla_dump")
out = run_py("""
import jax, jax.numpy as jnp
def simple_computation(x, y):
return jnp.exp(jnp.log(x) + y)
x = jnp.ones((1000, 1000))
print(float(jax.jit(simple_computation)(x, x).sum()))
""", env_extra={"XLA_FLAGS": f"--xla_dump_to={dump_dir} --xla_dump_hlo_as_text"})
dumped = sorted(os.listdir(dump_dir))
mine = [f for f in dumped if "jit_simple_computation" in f]
print(f"dumped {len(dumped)} files in total; the ones for jit_simple_computation:")
print("\n".join(" " + f for f in mine))
before = open(os.path.join(dump_dir, [f for f in mine if f.endswith("before_optimizations.txt")][0])).read()
after = open(os.path.join(dump_dir, [f for f in mine if f.endswith("cpu_after_optimizations.txt")][0])).read()
print("\n---- before_optimizations (ENTRY part) ----\n", clean_hlo(before).split("ENTRY")[-1])
print("---- after_optimizations (ENTRY part) ----\n", clean_hlo(after).split("ENTRY")[-1])
assert count_fusions(before) == 0 and count_fusions(after) == 1
Case study: why inspect the compiled output — a CPU backend bug that silently returns wrong results¶
The course says "inspecting the jaxpr only shows what's sent to XLA, not the result of XLA's optimizations". The real case below shows that sometimes you have to look at the optimized HLO to pin down a problem.
The CPU backend of jaxlib 0.12.0.1 hands patterns like matmul followed by reduce to the YNNPACK library, which shows up in HLO as a fusion with kind=kCustom and "kind":"__ynn_fusion". Testing found that when a constant matrix captured by a closure is the left operand of a matmul followed by a sum, the result silently comes out wrong. No error, but the numbers are wrong.
Separate tests on this machine with square $n\times n$ matrices showed: with the constant as the right operand (x @ W), or without a reduce after the multiply, the results are correct. The size threshold that triggers it depends on the shape; with this set of random numbers it goes wrong from $n=64$ on. Version bisection, whether float64 is affected, narrower workaround flags, etc. are collected in "Known issues in this environment" in Basics 00 Overview; this section only does the root-cause analysis.
rng = np.random.default_rng(0)
n = 128
W_np = rng.standard_normal((n, n)).astype(np.float32)
X_np = rng.standard_normal((n, n)).astype(np.float32)
ref = float(np.sum(W_np.astype(np.float64) @ X_np.astype(np.float64))) # float64 reference
W_const = jnp.asarray(W_np)
closure_fn = jax.jit(lambda x: jnp.sum(W_const @ x)) # W as a closure constant
param_fn = jax.jit(lambda x, w: jnp.sum(w @ x)) # W passed as an argument
r_closure = float(closure_fn(jnp.asarray(X_np)))
r_param = float(param_fn(jnp.asarray(X_np), W_const))
print(f"float64 reference = {ref:.4f}")
print(f"argument version = {r_param:.4f} error {abs(r_param - ref):.2e}")
print(f"closure version = {r_closure:.4f} error {abs(r_closure - ref):.2e}")
print("did the closure version trigger the bug:", abs(r_closure - ref) > 1.0)
assert abs(r_param - ref) < 1e-2 # the argument version can always be trusted
h_closure = clean_hlo(closure_fn.lower(jnp.asarray(X_np)).compile().as_text())
h_param = clean_hlo(param_fn.lower(jnp.asarray(X_np), W_const).compile().as_text())
def fused_body(h):
body = h.split("%fused_computation")[1]
return body[:body.index("\n}")]
print("==== fused_computation of the closure version ====", fused_body(h_closure), sep="")
print("==== fused_computation of the argument version ====", fused_body(h_param), sep="")
print("ENTRY instruction:\n ", [l.strip() for l in h_closure.splitlines() if "ROOT %ynn" in l][0][:120])
assert '"kind":"__ynn_fusion"' in h_closure and '"kind":"__ynn_fusion"' in h_param
# The key difference: the closure version embeds the 128×128 constant directly inside the YNN fusion, as the left operand of dot
assert re.search(r"constant\.\d+ = f32\[128,128\]\{1,0\} constant\(\{\.\.\.\}\)", fused_body(h_closure))
assert "constant({...})" not in fused_body(h_param)
The two HLO graphs are exactly the same except for one thing: in the closure version, the 128×128 weight is embedded inside __ynn_fusion as constant({...}) and used as the left operand of dot; in the argument version it's a parameter. We can infer that the problem lies in YNNPACK's path for "embedded constant + dot + reduce", not in JAX's tracing stage; nothing looks wrong in the jaxpr.
To check this inference, recompile the same closure function in a subprocess with XLA's debug option --xla_cpu_experimental_ynn_fusion_type= (set to an empty list, i.e. YNN library fusion turned off):
probe_src = """
import sys, numpy as np, jax, jax.numpy as jnp
rng = np.random.default_rng(0); n = 128
W = rng.standard_normal((n, n)).astype(np.float32); X = rng.standard_normal((n, n)).astype(np.float32)
Wc = jnp.asarray(W)
f = jax.jit(lambda x: jnp.sum(Wc @ x))
hlo = f.lower(jnp.asarray(X)).compile().as_text()
print(float(f(jnp.asarray(X))), float(np.sum(W.astype(np.float64) @ X.astype(np.float64))), "__ynn_fusion" in hlo)
"""
res_default = run_py(probe_src).split()
res_noynn = run_py(probe_src, env_extra={"XLA_FLAGS": "--xla_cpu_experimental_ynn_fusion_type="}).split()
print("default settings : result %s reference %s uses YNN=%s" % tuple(res_default[-3:]))
print("YNN fusion off : result %s reference %s uses YNN=%s" % tuple(res_noynn[-3:]))
val, refv, ynn = float(res_noynn[-3]), float(res_noynn[-2]), res_noynn[-1]
assert ynn == "False" and abs(val - refv) < 1e-2 # with YNN fusion off, the closure version is correct too
Conclusions and engineering lessons:
- Pass large arrays as arguments; don't capture them in closures. That sidesteps this bug and also avoids embedding large constants in the executable, which slows compilation (covered in the "global variables" part of Basics 02 JIT compilation).
- A normal-looking jaxpr doesn't mean the result is correct. When numbers look suspicious, compare the optimized HLO (
clean_hloplus a diff), and cross-check against NumPy in float64. - The debug options in
XLA_FLAGS(--xla_dump_to,--xla_disable_hlo_passes=..., backend-specific--xla_cpu_*) are great tools for bisecting compiler problems, but they only take effect at process startup.
What this means for JAX developers (course key points)¶
- Writing code with vectorized
jax.numpyoperations usually gives fusible op sequences that XLA can optimize well; - Compared with jitting many small functions separately, wrapping a larger, whole function in
@jitusually gives XLA more room to optimize; - Inspecting the jaxpr shows the ops sent to XLA, but not the optimization result; for that, look at the optimized HLO;
- With XLA, high-level numerical programs written in Python can approach the performance of hand-written C++/CUDA.
🔍 Extra¶
⚪ Optional: useful when you read CPU HLO or compare implementations with cost_analysis.
- The CPU backend doesn't rely on LLVM alone. The jaxlib 0.12 CPU backend hands dot, reduce and some elementwise chains to the YNNPACK library (shown in HLO as
kind=kCustomplus__ynn_fusion), while ordinary elementwise fusions arekind=kLoopwith code generated by XLA itself. Don't be surprised to seeynn_fusionwhen reading CPU HLO and traces. lowered.as_text(debug_info=True)includes source locations; themetadata={op_name=... stack_frame_id=...}incompiled.as_text()lets you trace each HLO instruction back to a line of Python.- The
flopsfromcost_analysis()is an estimate: $2MNK$ fordot, 1 per element for elementwise ops, withtranscendentalscounted separately. It's handy for comparing the "theoretical work" of different implementations (used in 2.8).
2.4 Memory layout and its effect on performance¶
A program's performance depends not only on which operations it does, but also on how efficiently data is read from memory. The way a multidimensional array is arranged in linear memory is called its memory layout. On accelerators with dedicated memory access hardware, such as GPUs and TPUs, layout matters especially.
Row-major and column-major¶
Take the 2×3 matrix $A=\begin{pmatrix}1&2&3\\4&5&6\end{pmatrix}$:
- Row-major (C style): linear memory holds
1, 2, 3, 4, 5, 6in order. This is NumPy's default layout, and JAX usually uses it too. - Column-major (Fortran style):
1, 4, 2, 5, 3, 6in order.
For an $m\times n$ matrix, the linear offset of element $(i,j)$ is:
$$ \text{offset}_{\text{row-major}}(i,j)= i\cdot n + j,\qquad \text{offset}_{\text{col-major}}(i,j)= i + j\cdot m . $$
In higher dimensions: when walking through contiguous memory, row-major varies the last index fastest, and column-major varies the first index fastest.
A = np.array([[1, 2, 3], [4, 5, 6]], dtype=np.float32)
A_F = np.asfortranarray(A)
print("row-major linear memory:", A.ravel(order="K"), " strides =", A.strides)
print("column-major linear memory:", A_F.ravel(order="K"), " strides =", A_F.strides)
assert A.ravel(order="K").tolist() == [1, 2, 3, 4, 5, 6]
assert A_F.ravel(order="K").tolist() == [1, 4, 2, 5, 3, 6]
# Offset formulas (in elements)
m_, n_ = A.shape
for i in range(m_):
for j_ in range(n_):
assert A.strides[0] // 4 * i + A.strides[1] // 4 * j_ == i * n_ + j_
assert A_F.strides[0] // 4 * i + A_F.strides[1] // 4 * j_ == i + j_ * m_
Why layout matters to hardware¶
- CPU caches: data is loaded in cache lines (commonly 64 bytes). Sequential access exploits spatial locality and gets a high cache hit rate. Row-major suits row-wise traversal, column-major suits column-wise traversal.
- GPU memory coalescing: if the threads of a warp access contiguous locations in global memory, those accesses can be merged into one wide transaction; strided access has to be split into many transactions, and effective bandwidth drops sharply. In row-major order, access along the last dimension (
A[i,j],A[i,j+1], …) usually coalesces, while access along a column (A[i,j],A[i+1,j], …) is usually strided. - TPU: high-bandwidth memory feeds the MXU, and it's most efficient when data is arranged in specific tiled formats. XLA's TPU backend does layout conversions to suit the MXU, but the initial layout decides how much extra cost those conversions incur.
Layouts in JAX and XLA¶
JAX arrays are row-major by default, same as NumPy. But inside a jitted function, XLA takes over. XLA does layout assignment and may switch an intermediate array to a layout that suits later operations such as a matmul (e.g. by doing a transpose), even if the input is row-major. So you usually can't explicitly control the exact layout of every intermediate array inside a jitted function, and you usually don't need to.
In HLO, the layout is written in braces after the type, e.g. f32[2000,2000]{1,0}. What's listed is the minor-to-major order: {1,0} means dimension 1 varies fastest, i.e. row-major.
M2 = jax.random.normal(jax.random.key(1), (2000, 2000))
print("jax.Array.format:", M2.format.layout)
h_id = compiled_hlo(lambda t: t * 2.0, M2)
print("HLO entry layout:", re.search(r"entry_computation_layout=\{([^}]*\}[^}]*)\}", h_id).group(1))
assert M2.format.layout.major_to_minor == (0, 1) # row-major: dimension 0 is the "slowest"
assert "f32[2000,2000]{1,0}" in h_id # HLO's minor_to_major notation
How XLA handles jnp.transpose¶
A transpose in a jaxpr is only a logical operation; XLA decides whether to really move data based on how it's used. Let's measure four cases:
B2 = jax.random.normal(jax.random.key(2), (2000, 2000))
cases = {
"① transpose only a.T": (lambda a: a.T, (M2,)),
"② transpose then elementwise a.T*2+1": (lambda a: a.T * 2 + 1, (M2,)),
"③ transpose then matmul a.T @ b": (lambda a, b: a.T @ b, (M2, B2)),
"④ double transpose a.T.T": (lambda a: a.T.T, (M2,)),
}
hlos = {}
for name, (fn, args) in cases.items():
jp = jax.make_jaxpr(fn)(*args)
h = clean_hlo(compiled_hlo(fn, *args))
hlos[name] = h
entry = [l.strip()[:110] for l in h.split("ENTRY")[-1].splitlines() if "ROOT" in l][0]
print(f"{name}\n transposes in the jaxpr: {count_primitives(jp.jaxpr)['transpose']}\n optimized ROOT: {entry}")
assert " transpose(" in hlos["① transpose only a.T"] # really moves data (inside a fusion)
assert count_fusions(hlos["② transpose then elementwise a.T*2+1"]) == 1 # transpose and elementwise fused into one kernel
assert " transpose(" not in hlos["③ transpose then matmul a.T @ b"] # transpose absorbed into dot's contracting dims
assert "lhs_contracting_dims={0}" in hlos["③ transpose then matmul a.T @ b"]
assert " transpose(" not in hlos["④ double transpose a.T.T"] # the two transposes cancel out
A few takeaways:
- ① A standalone
a.Tmust produce a new row-major array, so there's one real data movement (transposepluscopy, fused into one kernel); - ② When a transpose is followed by elementwise ops, the transpose is fused into the same kernel, which simply reads with transposed indices; there's no separate extra round trip to memory;
- ③
a.T @ bhas no transpose op at all: XLA just changesdot's contracting dimension tolhs_contracting_dims={0}, and the matmul kernel reads the data in transposed order; - ④
a.T.Tis simplified away, leaving only onecopy. The copy is still needed because the output can't share a buffer with the input unless you usedonate_argnums.
So inside jit, you usually don't need to worry about the cost of transposes when writing x.T @ y or jnp.einsum. What you really need to care about are materialized transposes (①) and strided access.
Timing contiguous vs strided access (CPU, printed only, no conclusions drawn)¶
sum_rows = jax.jit(lambda a: a.sum(axis=1)) # along the last dimension: contiguous access
sum_cols = jax.jit(lambda a: a.sum(axis=0)) # along the first dimension: strided access (XLA may reorder loops)
copy_c = jax.jit(lambda a: a * 1.0 + 0.0) # plain copy (simplified to copy)
copy_t = jax.jit(lambda a: a.T) # materialized transpose
t_rows, t_cols = bench(sum_rows, M2), bench(sum_cols, M2)
t_copy, t_tr = bench(copy_c, M2), bench(copy_t, M2)
print(f"XLA sum(axis=1) {t_rows*1e3:7.3f} ms | sum(axis=0) {t_cols*1e3:7.3f} ms")
print(f"XLA copy {t_copy*1e3:7.3f} ms | materialized transpose {t_tr*1e3:7.3f} ms")
Mn = np.asarray(M2)
t0 = time.perf_counter(); [np.ascontiguousarray(Mn) .copy() for _ in range(10)]; t_np_c = (time.perf_counter()-t0)/10
t0 = time.perf_counter(); [np.ascontiguousarray(Mn.T) for _ in range(10)]; t_np_t = (time.perf_counter()-t0)/10
print(f"NumPy contiguous copy {t_np_c*1e3:7.3f} ms | transpose made contiguous {t_np_t*1e3:7.3f} ms")
print("(XLA picks a better traversal order for reduce, so the gap between the two sums is often smaller than with a naive C loop)")
Practical considerations (the course's five points)¶
- Identify memory bottlenecks: first use the profiler to check whether you're memory-bound. If operations on large arrays are slow, the access pattern may be poor, which relates to layout.
- Understand the hardware: on GPU, remember that access along the last dimension (row-major) is what coalesces; on TPU, certain shapes and contracting dimensions of
jnp.einsumandjnp.matmulmap to the hardware more easily. - Algorithms and data organization: if you frequently access columns on GPU and profiling confirms that's the bottleneck, consider transposing once up front, but count the cost of the transpose itself.
- Large data transfers: NumPy can represent column-major arrays; before handing one to JAX, it's best to make sure it's standard row-major, to avoid paying a conversion cost right at the start. We measure this below.
- Custom operations: when writing custom primitives or integrating external code (Chapter 5 of the course), you're responsible for handling layouts on the target backend yourself.
big_c = np.random.default_rng(0).random((3000, 3000), dtype=np.float32)
big_f = np.asfortranarray(big_c)
t_put_c = bench(lambda a: jax.device_put(a), big_c, n=10)
t_put_f = bench(lambda a: jax.device_put(a), big_f, n=10)
print(f"device_put row-major {t_put_c*1e3:8.3f} ms | column-major {t_put_f*1e3:8.3f} ms (about {t_put_f/t_put_c:.0f}×)")
arr_from_f = jnp.asarray(big_f)
assert arr_from_f.format.layout.major_to_minor == (0, 1) # once inside JAX, everything is row-major
assert np.array_equal(np.asarray(arr_from_f), big_c) # values unchanged
🔍 Extra¶
⚪ Optional: read this when your data pipeline moves lots of data from NumPy, or when you need to control input/output layouts.
- On CPU,
device_putof a C-contiguous, well-aligned NumPy array can be zero-copy (it takes only microseconds). A Fortran-order array has to be converted to row-major first, which needs a full copy. The measurement above shows a gap of about two orders of magnitude (the exact factor varies with the machine and array size). In a data pipeline, normalizing arrays up front withnp.ascontiguousarrayis one of the cheapest optimizations. x.format.layoutshows the physical layout of ajax.Array.jax.jit'sin_shardings/out_shardingsalso accept aFormat(from jax.experimental.layout import Format, Layout, e.g.Format(Layout.AUTO)), which lets XLA choose input/output layouts automatically (advanced usage, mainly beneficial on TPU).- To judge whether a transpose is "expensive", look for a materialized
transposein the optimized HLO, not in the jaxpr.
2.5 Avoiding recompilation¶
jit speeds things up by compiling functions with XLA, but compilation itself has a cost. JAX first traces the Python function to produce a jaxpr; XLA then compiles the jaxpr into optimized code for the target hardware and for the input shapes and types seen at tracing time.
Caching: the first time a @jit function is called with some shape/type, it traces and compiles, and the resulting executable goes into a cache. Later calls with the same shapes and types reuse the cache with no tracing or compilation cost. If the characteristics of the arguments change so that a new jaxpr is needed, it has to retrace and recompile. If that happens often, say in a loop, the benefit of jit is wiped out.
The course lists six causes of recompilation (for the details, trust what the code in this section actually shows):
- Argument shape: first (10,5), then (20,5);
- Argument dtype: e.g. switching from float32 to float64 (see the Note below);
- Python type of an argument: passing an int at one point and a float at another in the same position;
- Value-dependent Python control flow (see 2.2: on traced values it errors right away; only static arguments get as far as recompilation);
- Functions as arguments: a different function object on each call;
- Closures and globals: when their values affect the graph structure (e.g. determine the size of an intermediate array).
For the full table of recompilation triggers and the basics of jax_log_compiles and _cache_size(), see "Table of recompilation triggers" in Basics 02 JIT compilation; for the complete diagnose-and-fix workflow, see topic 1 in Basics 07 Practical topics. This section verifies the course's three code blocks and adds material the Basics part didn't cover.
The first code block (block A) tries to show "recompilation when the value of size_param changes", but it does not mark size_param as static:
# ❌ Old style (buggy code): a Python if on a traced argument can't be traced, so the very first call errors; "recompiling on the third call" never comes into it
@jax.jit
def example_function(x, size_param):
# Python control flow depends on the value of size_param
if size_param > 5:
return jnp.sum(x * 2.0)
else:
return jnp.sum(x + 1.0)
data = jax.random.normal(jax.random.key(0), (10,))
try:
example_function(data, 3).block_until_ready()
err_a = None
except Exception as e:
err_a = type(e).__name__
print("measured:", err_a)
assert err_a == "TracerBoolConversionError"
⚠️ Note: Block A will not show "the first and third calls are noticeably slower, the third because the value of
size_paramchanged the control-flow path". Its very first call raisesTracerBoolConversionError, because inside jit ordinary arguments are tracers and can't be used in a Pythonif. To see recompilation by value, you need block B'sstatic_argnums.
Strategy 1: static arguments (block B)¶
When an argument's value affects the graph structure (shapes, hyperparameters, control flow), mark it static with static_argnums / static_argnames, and JAX compiles a separate version for each combination of static values. Note that if a static argument takes too many values, you end up compiling a pile of specialized versions, which adds compile time and cache usage. So it's only suitable for arguments that take a limited set of values and determine the graph structure.
# ✅ Recommended (course block B): mark it static; use CompileCounter instead of "guessing compilation from timings"
@jax.jit(static_argnums=(1,)) # mark size_param (index 1) as static (0.12 supports this decorator-factory form)
def example_function_static(x, size_param):
print(f"Compiling for size_param = {size_param}") # runs only during tracing
if size_param > 5:
return jnp.sum(x * 2.0)
else:
return jnp.sum(x + 1.0)
timings = []
with CompileCounter("example_function_static") as cc:
for sp in [3, 3, 7, 7]:
t0 = time.perf_counter()
example_function_static(data, sp).block_until_ready()
timings.append(time.perf_counter() - t0)
for sp, t in zip([3, 3, 7, 7], timings):
print(f"size_param={sp}: {t*1e3:8.3f} ms")
print("actual compilations:", cc.count)
assert cc.count == 2 # 3 and 7 each compile once
Strategy 2: keep input types and shapes consistent¶
Variable-length sequences can be padded to a fixed maximum length, with a mask to ignore the padding (Basics 02 JIT compilation has a practical bucketed-padding example). When types don't match, convert the dtype explicitly before calling the jitted function.
⚠️ Note: It's easy to think that "passing float64 after float32 can also trigger recompilation". Under the default
jax_enable_x64=False, when you pass a NumPy float64 array, JAX first canonicalizes it to float32, so it matches the existing float32 cache entry and does not recompile. Only with x64 enabled is float64 a separate type. The real trap is weak types: the Python scalars1,1.0andjnp.float32(1)are three different cache keys.
@jax.jit
def scale(v):
return v * 2
with CompileCounter("scale") as cc:
scale(jnp.ones(3, jnp.float32))
scale(np.ones(3, np.float64)) # canonicalized to float32 in x32 mode
n_after_f64 = cc.count
scale(1) # int32, weak_type=True
scale(1.0) # float32, weak_type=True
scale(jnp.float32(1.0)) # float32, weak_type=False
for r in cc.records:
print(" ", r.split(" with global shapes and types ")[1].split(". Argument")[0])
print("float32 array + float64 NumPy array → compilations:", n_after_f64, "; total compilations:", cc.count)
assert n_after_f64 == 1 and cc.count == 4
assert not jax.config.jax_enable_x64
Strategy 3: use JAX control-flow primitives (block C)¶
Replace Python if/for/while that depend on traced values with lax.cond, lax.scan and lax.while_loop. The branches or loops get embedded in the XLA computation graph, and value changes don't require recompilation.
@jax.jit
def example_function_lax(x, size_param_val):
# size_param_val is a 0-d array (or a traceable Python scalar)
pred = size_param_val > 5
def true_fun(operand):
return jnp.sum(operand * 2.0)
def false_fun(operand):
return jnp.sum(operand + 1.0)
return lax.cond(pred, true_fun, false_fun, x)
size_val_3, size_val_7 = jnp.array(3), jnp.array(7)
with CompileCounter("example_function_lax") as cc:
r3 = example_function_lax(data, size_val_3).block_until_ready()
r7 = example_function_lax(data, size_val_7).block_until_ready()
print("compilations:", cc.count, "| r3 =", r3, "| r7 =", r7)
assert cc.count == 1
np.testing.assert_allclose(r3, np.sum(np.asarray(data) + 1.0), rtol=1e-5)
np.testing.assert_allclose(r7, np.sum(np.asarray(data) * 2.0), rtol=1e-5)
Strategy 4: manage closures and function arguments carefully¶
- Inside a jitted function, don't refer to outside variables that change often and affect the graph structure; turn them into arguments (marked static if necessary);
- When passing a function, pass the same function object every time;
- Define functions at the top level where possible, and avoid defining functions dynamically inside loops.
People often say "functools.partial sometimes helps you get a stable function object". That statement needs to be more precise:
def base(v, a):
return v * a
v3 = jnp.ones(3)
with CompileCounter("<lambda>") as c_lambda:
for _ in range(3):
jax.jit(lambda v: v * 2)(v3) # a new lambda object every time
with CompileCounter("base") as c_partial:
for _ in range(3):
jax.jit(functools.partial(base, a=2.0))(v3) # a new partial object every time
with CompileCounter("base") as c_rejit:
for _ in range(3):
jax.jit(base)(v3, 2.0) # a new jit wrapper, but the wrapped function is the same
p_once = jax.jit(functools.partial(base, a=2.0)) # the partial is created only once
with CompileCounter("base") as c_once:
for _ in range(3):
p_once(v3)
print(f"new lambda in the loop: {c_lambda.count} | new partial in the loop: {c_partial.count} | "
f"jax.jit(base) in the loop: {c_rejit.count} | partial created once: {c_once.count} compilations")
assert (c_lambda.count, c_partial.count, c_rejit.count, c_once.count) == (3, 3, 1, 1)
⚠️ Note:
functools.partialcreates a new object every time it's called, and it doesn't define equality by content. So writingjax.jit(partial(f, a=...))in a loop recompiles every time (3 calls, 3 compilations), just like creating a new lambda each time. The right approach is to create the partial or jit wrapper once and reuse it, or pass the bound value as a (static) argument. Interestingly,jax.jit(base)hits the cache even when it's recreated over and over in a loop, because the cache is indexed by the underlying function object.
🔍 Extra: jax_explain_cache_misses: let JAX tell you why the cache missed¶
⚪ Optional: read this when you're tracking down "why did it compile again"; the full diagnostic workflow is in Basics 07 topic 1.
jax_log_compiles only tells you "it compiled". jax_explain_cache_misses goes further: it compares against the closest old cache key and points out exactly which property of which argument changed. Here we demo only a shape change; for weak types, pytree structure changes, and the case where a change in a static argument's value only prints explanation unavailable, see topic 1.4 in Basics 07 Practical topics and §5.2 of Advanced 08 Debugging toolchain.
@jax.jit
def explain_me(v, w):
return v @ w
class _Grab(logging.Handler):
def __init__(self):
super().__init__(); self.msgs = []
def emit(self, record):
self.msgs.append(record.getMessage())
grab = _Grab(); jl = logging.getLogger("jax"); jl.addHandler(grab)
saved = [(h, h.level) for h in jl.handlers if h is not grab]
for h, _ in saved: h.setLevel(logging.ERROR) # don't flood the output, just collect
jax.config.update("jax_explain_cache_misses", True)
try:
explain_me(jnp.ones((4, 4)), jnp.ones((4, 4)))
explain_me(jnp.ones((4, 4)), jnp.ones((4, 8))) # only w's shape changed
finally:
jax.config.update("jax_explain_cache_misses", False)
for h, lvl in saved: h.setLevel(lvl)
jl.removeHandler(grab)
miss = [m for m in grab.msgs if "CACHE MISS" in m and "explain_me" in m]
print(miss[-1])
assert "at w, now f32[4,8] and before f32[4,4]" in miss[-1]
🔍 Extra: the persistent compilation cache: so a second process doesn't recompile either¶
⚪ Optional: only useful for large models where compilation takes minutes and you restart processes often.
The caches discussed so far live in memory and vanish when the process exits. For large models, every script restart means recompiling, which can take minutes. The persistent compilation cache writes XLA executables to disk, and later processes load them directly as long as the cache key matches.
| Setting | What it does | Default |
|---|---|---|
jax_compilation_cache_dir | cache directory (the environment variable JAX_COMPILATION_CACHE_DIR also works) | unset, i.e. disabled |
jax_persistent_cache_min_compile_time_secs | only write to disk if compilation took longer than this | 1.0 seconds, so small functions aren't cached by default |
jax_persistent_cache_min_entry_size_bytes | minimum entry size (-1 means no limit) | 0 |
The cache key includes the HLO module contents, the jax/jaxlib versions, XLA compile options (including XLA_FLAGS), the device type and so on, so upgrading JAX or changing flags automatically invalidates old entries.
Below we verify it with two independent subprocesses sharing one cache directory: the first process MISSes and writes to disk, the second one HITs.
cache_dir = os.path.join(SCRATCH, "jax_cache")
cache_src = """
import sys, time, logging
import jax, jax.numpy as jnp
jax.config.update("jax_compilation_cache_dir", sys.argv[1])
jax.config.update("jax_persistent_cache_min_compile_time_secs", 0.0) # for the demo: write even small functions to disk
jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
logging.basicConfig(level=logging.WARNING, format="%(message)s")
logging.getLogger("jax._src.compiler").setLevel(logging.DEBUG) # print HIT / MISS
def model(x):
for _ in range(20):
x = jnp.tanh(x @ x) + 0.1
return x
t0 = time.perf_counter()
exe = jax.jit(model).lower(jnp.ones((64, 64))).compile()
print("COMPILE_SECONDS", time.perf_counter() - t0)
"""
out1 = run_py(cache_src, args=(cache_dir,))
out2 = run_py(cache_src, args=(cache_dir,))
pick = lambda s: [l for l in s.splitlines() if "jit_model" in l or "COMPILE_SECONDS" in l]
print("---- process 1 ----", *pick(out1), sep="\n")
print("---- process 2 ----", *pick(out2), sep="\n")
print("\ncache directory contents:", sorted(f[:40] + "…" for f in os.listdir(cache_dir) if f.endswith("-cache")))
assert "PERSISTENT COMPILATION CACHE MISS for 'jit_model'" in out1
assert "Persistent compilation cache hit for 'jit_model'" in out2
assert any(f.startswith("jit_model-") for f in os.listdir(cache_dir))
Key points:
jax.config.update("jax_compilation_cache_dir", ...)must be called before the first compilation, ideally at the very top of the script;- The default 1-second threshold is deliberate: for small functions, reading from disk can cost more than recompiling;
- Several processes or machines can share one directory (e.g. a network file system or a GCS path), which is common in cluster training;
- The persistent cache only saves XLA compilation; tracing and lowering still run, so it doesn't help with slow tracing at the Python level (such as a huge unrolled loop).
2.6 Operator fusion and op optimization¶
One of the most important optimizations XLA performs at compile time is operator fusion: merging several operations into one larger compute kernel. Understanding it helps you read profiling results and avoid accidentally preventing this optimization.
What fusion is: at its core, it merges a chain of elementwise operations with producer–consumer relationships into one composite operation. The course example:
# ❌ Old style (buggy code): x and y generated with the same PRNGKey → the two arrays are identical
key_old = jax.random.PRNGKey(0)
x_old = jax.random.uniform(key_old, (1000, 1000))
y_old = jax.random.uniform(key_old, (1000, 1000))
print("x and y are identical:", bool(jnp.array_equal(x_old, y_old)))
assert bool(jnp.array_equal(x_old, y_old))
# ✅ Recommended: jax.random.key + split
import jax
import jax.numpy as jnp
def simple_computation(x, y):
a = jnp.log(x)
b = a + y
c = jnp.exp(b)
return c
# JIT-compile the function
compiled_computation = jax.jit(simple_computation)
# Example data
kx, ky = jax.random.split(jax.random.key(0))
x = jax.random.uniform(kx, (1000, 1000))
y = jax.random.uniform(ky, (1000, 1000))
assert not bool(jnp.array_equal(x, y))
# Run
result = compiled_computation(x, y).block_until_ready()
np.testing.assert_allclose(result, np.exp(np.log(np.asarray(x, np.float64)) + np.asarray(y, np.float64)), rtol=1e-5)
print("result is correct, shape", result.shape)
Without fusion (e.g. running ops one by one on a GPU), simple_computation launches three kernels:
- Read
xfrom memory, computelog(x), write the result back toa; - Read
aandy, computea + y, write the result back tob; - Read
b, computeexp(b), write the result back toc.
Each step reads and writes the accelerator's main memory (e.g. GPU HBM), and that memory traffic is often the real bottleneck. When XLA analyzes the HLO (XLA doesn't read jaxprs directly; see the compilation pipeline in 2.3), it notices that the intermediates a and b are used only by the next op, so it merges the three steps into one kernel: load x and y once, compute exp(log(x) + y) in on-chip storage (registers/cache), and write c back just once.
With $N$ float32 elements per array, a rough estimate of main-memory traffic:
$$ \underbrace{(1+1) + (2+1) + (1+1)}_{\text{unfused: reads+writes}} = 7\times 4N \text{ bytes} \qquad\text{vs.}\qquad \underbrace{2 + 1}_{\text{fused}} = 3\times 4N \text{ bytes}. $$
Benefits of fusion:
- Less memory bandwidth (the biggest one): accessing main memory is orders of magnitude slower than registers/on-chip cache;
- Lower kernel launch overhead: in large computations this usually matters less than bandwidth, but for functions made of many small ops it can be significant;
- More instruction-level parallelism: bigger kernels give the hardware scheduler more independent instructions to arrange.
How to observe fusion (the course's three methods): ① in the profiler, the number of kernels on the accelerator timeline is smaller than the number of ops in the source; ② the @jit version is clearly faster than the sum of running each op separately; ③ inspect the XLA HLO. Below we go straight to method ③: count fusions in the optimized HLO, and compare against turning off the fusion pass with --xla_disable_hlo_passes=fusion (again, this has to be done in a subprocess):
h_fused = compiled_hlo(simple_computation, x, y)
print("default (fused):\n", clean_hlo(h_fused).split("ENTRY")[-1])
ca_fused = jax.jit(simple_computation).lower(x, y).compile().cost_analysis()
fusion_probe = """
import json, re, jax, jax.numpy as jnp, time
def simple_computation(x, y):
return jnp.exp(jnp.log(x) + y)
kx, ky = jax.random.split(jax.random.key(0))
x = jax.random.uniform(kx, (1000, 1000)); y = jax.random.uniform(ky, (1000, 1000))
c = jax.jit(simple_computation).lower(x, y).compile()
h = c.as_text()
f = jax.jit(simple_computation); f(x, y).block_until_ready()
t0 = time.perf_counter()
for _ in range(50): f(x, y).block_until_ready()
entry = [l.strip().split(", metadata")[0] for l in h.split("ENTRY")[-1].splitlines() if " fusion(" in l]
print("JSON" + json.dumps(dict(n_fusion=len(re.findall(r" fusion\\(", h)), entry=entry,
bytes=c.cost_analysis()["bytes accessed"],
temp=c.memory_analysis().temp_size_in_bytes,
t=(time.perf_counter() - t0) / 50)))
"""
def parse_json(out):
return json.loads([l for l in out.splitlines() if l.startswith("JSON")][-1][4:])
r_on = parse_json(run_py(fusion_probe))
r_off = parse_json(run_py(fusion_probe, env_extra={"XLA_FLAGS": "--xla_disable_hlo_passes=fusion"}))
print("ENTRY instructions with the fusion pass turned off:")
for l in r_off["entry"]:
print(" ", l[:120])
print(f"\n{'':10s}{'fusions':>10s}{'bytes accessed':>18s}{'temp buffer':>14s}{'time (ms)':>10s}")
for name, r in [("fused", r_on), ("unfused", r_off)]:
print(f"{name:10s}{r['n_fusion']:>10d}{r['bytes']:>18,.0f}{r['temp']:>14,d}{r['t']*1e3:>10.3f}")
N_el = 1000 * 1000
assert count_fusions(h_fused) == 1 and r_on["n_fusion"] == 1 and r_off["n_fusion"] == 3
assert r_on["bytes"] == 3 * 4 * N_el and r_off["bytes"] == 7 * 4 * N_el # matches the 3N vs 7N estimate above exactly
XLA's own cost_analysis reports bytes accessed: $3\times4N$ fused and $7\times4N$ unfused, exactly matching the hand calculation above. One more detail: unfused, temp is still 0. That's because the CPU backend's buffer assignment lets elementwise ops write in place into the output buffer (a and b both reuse c's memory). In other words, the cost of not fusing is mainly memory traffic (an extra 16 MB read and written), not necessarily memory footprint.
Next, using the course's method ②, compare eager (op-by-op execution) with jit:
⚠️ Note: Without
@jit, intermediate results are not "materialized as full NumPy arrays". In eager mode, eachjnpop is compiled and dispatched on its own, and the intermediates arejax.Arrays that stay on the device; they don't go back to the host. Eager is slow because every op reads and writes memory in full, plus the dispatch overhead each time.
a_eager = jnp.log(x)
print("type of the eager intermediate:", type(a_eager).__name__, "| device:", a_eager.device)
assert isinstance(a_eager, jax.Array) and not isinstance(a_eager, np.ndarray)
t_eager = bench(simple_computation, x, y)
t_jit = bench(compiled_computation, x, y)
print(f"eager (3 dispatches, 3 kernels): {t_eager*1e3:.3f} ms | jit (1 fused kernel): {t_jit*1e3:.3f} ms | ratio {t_eager/t_jit:.2f}×")
Tips for encouraging fusion (from the course)¶
- Keep elementwise ops together: a chain of
jax.numpyelementwise ops fuses easily; - Avoid "fusion barriers": forced synchronization, complex data dependencies XLA can't analyze, some control flow, and explicit memory layout changes between fusible ops can all prevent adjacent ops from fusing;
- Trust the compiler: write clear, numerically correct code with standard library functions. Trying to "force fusion" by hand usually backfires.
Below we look at two common "boundaries": the explicit barrier lax.optimization_barrier, and a reduction followed by a broadcast (the max/sum in softmax):
mat = jax.random.normal(jax.random.key(3), (1000, 1000))
f_chain = lambda t: jnp.exp(jnp.log(jnp.abs(t) + 1.0) * 2.0)
f_barrier = lambda t: jnp.exp(lax.optimization_barrier(jnp.log(jnp.abs(t) + 1.0)) * 2.0)
def softmax_manual(t):
m = t.max(axis=-1, keepdims=True)
e = jnp.exp(t - m)
return e / e.sum(axis=-1, keepdims=True)
for name, fn in [("elementwise chain", f_chain), ("with optimization_barrier", f_barrier),
("hand-written softmax", softmax_manual), ("jax.nn.softmax", jax.nn.softmax)]:
h = compiled_hlo(fn, mat)
kinds = re.findall(r'kind=(\w+)', h)
print(f"{name:28s} fusions = {count_fusions(h)} kinds = {kinds}")
assert count_fusions(compiled_hlo(f_chain, mat)) == 1
assert count_fusions(compiled_hlo(f_barrier, mat)) == 2 # the barrier splits one kernel into two
assert count_fusions(compiled_hlo(softmax_manual, mat)) > 1 # reduce → broadcast forms a kernel boundary
np.testing.assert_allclose(f_barrier(mat), f_chain(mat), rtol=1e-6)
np.testing.assert_allclose(softmax_manual(mat), jax.nn.softmax(mat), rtol=1e-5, atol=1e-7)
🔍 Extra¶
🟢 Must-read
- Kinds of fusion:
kLoopis a pure elementwise/broadcast fusion;kInputis a fusion built around a reduction (common on GPU);kOutputis a fusion built around a dot/conv with elementwise ops on the output side;kCustomis a fusion implemented by a library or a specialized code generator (__ynn_fusionon CPU, Triton GEMM fusions on GPU, etc.; on GPU, a matmul that calls cuBLAS directly shows up as acustom-calland doesn't count as a fusion). On CPU, the number of fusions you count is roughly the number of kernels you'll see in a trace. - Reductions are natural kernel boundaries: softmax needs to "compute the max of the whole row first, then process elementwise", so it's hard to make it a single kernel. That's exactly where hand-written kernels like FlashAttention earn their keep (covered in the Pallas chapter later in the course).
lax.optimization_barrieris meant to stop XLA from reordering or merging across it. It's commonly used to control rematerialization (remat) and for debugging; don't use it as a performance tool.- The most reliable evidence of whether fusion happened is
cost_analysis()["bytes accessed"](memory traffic), which is more reliable than timing, especially on CPU.memory_analysis().temp_size_in_bytesreflects "whether an intermediate got materialized" (Section 2.8 uses it for a controlled comparison), but elementwise intermediates can reuse the output buffer in place, so temp being 0 doesn't necessarily mean fusion happened.
2.7 Asynchronous dispatch¶
When you optimize JAX code on GPU/TPU, timings sometimes come out surprisingly fast. The reason is JAX's asynchronous dispatch.
Unlike standard Python, which waits for each step to finish, when JAX calls a function aimed at an accelerator (especially a jitted function) it does three things in order:
- Trace and compile (if needed): the first time it sees a given shape/type, or when recompilation is triggered;
- Dispatch: put the compiled kernel into the accelerator's execution queue;
- Return control immediately: hand control back to the Python interpreter without waiting for the computation on the device to finish.
The accelerator works through the queue in the background while Python keeps going. That way, control logic on the CPU and numerical computation on the accelerator can run in parallel. For example, while the GPU processes the current batch, the CPU can load and preprocess the next one, improving overall throughput.
Impact on benchmarking: if you time JAX calls directly with time.time() / time.perf_counter(), what you measure is often not the real execution time. The course's block A is the counterexample of "naive timing":
# ❌ Wrong: no warm-up and no block_until_ready (often described as "mostly measures dispatch time")
@jax.jit
def compute_heavy(m):
return jnp.dot(m, m.T)
x_h = jax.random.normal(jax.random.key(0), (2000, 2000)) # using the recommended key style; smaller than the original example's 4096×4096
with CompileCounter("compute_heavy") as cc:
start_time = time.perf_counter()
result = compute_heavy(x_h) # control returns almost immediately?
end_time = time.perf_counter()
naive = end_time - start_time
print(f"naive timing: {naive:.6f} s; compilations during this call = {cc.count}")
assert cc.count == 1 # the first call includes tracing + compilation, not just "dispatch"
# ✅ Recommended: warm up (compile) first, then time, with block_until_ready
result_compiled = compute_heavy(x_h).block_until_ready() # force compilation ahead of time
t0 = time.perf_counter()
result = compute_heavy(x_h)
t_dispatch = time.perf_counter() - t0 # only up to "dispatch done"
ready_right_after = result.is_ready()
result.block_until_ready() # wait for the computation to finish
t_total = time.perf_counter() - t0
print(f"first call (incl. compilation): {naive*1e3:9.3f} ms")
print(f"after warm-up, dispatch only : {t_dispatch*1e3:9.3f} ms (is_ready() right after dispatch = {ready_right_after})")
print(f"after warm-up, correct timing : {t_total*1e3:9.3f} ms")
print(f"share of the real time spent on dispatch: {t_dispatch / t_total:.1%}")
assert result.is_ready()
⚠️ Note: What block A measures is not "mostly the time to dispatch
jnp.dotto the accelerator". It has no warm-up, so the first call includes tracing and XLA compilation (verified above withCompileCounter: it compiled exactly once). So that number is neither dispatch time nor execution time. Only after warm-up does timing withoutblock_until_readyreflect just the dispatch time.
block_until_ready() can be called on any jax.Array; it blocks Python until the device computation producing that array finishes. There's also a function form, jax.block_until_ready(x):
# ❌ Common misconception: thinking that calling jax.block_until_ready() with no arguments waits for all pending computation on all devices
try:
jax.block_until_ready()
err_bur = None
except TypeError as e:
err_bur = str(e)
print("measured TypeError:", err_bur)
assert err_bur is not None and "missing 1 required positional argument" in err_bur
# ✅ Recommended: pass a pytree; it waits only for the jax.Array leaves in that pytree and returns the pytree as is
pending = {"a": compute_heavy(x_h), "b": [compute_heavy(x_h)[0]], "step": 7}
out_tree = jax.block_until_ready(pending)
assert out_tree is pending and out_tree["step"] == 7 # non-array leaves are kept as is
assert all(leaf.is_ready() for leaf in jax.tree.leaves(out_tree) if isinstance(leaf, jax.Array))
print("all arrays in the pytree are ready; the same object is returned:", out_tree is pending)
⚠️ Note:
jax.block_until_ready()doesn't wait for "all" pending asynchronous computation "on all devices". Calling it with no arguments raisesTypeError(missing the required argumentx). Called on a pytree, it waits only for the array leaves of that pytree, then returns the pytree as is; it doesn't wait for unrelated computation.
When to use block_until_ready() (from the course):
- Benchmarking: required when measuring JAX code running on an accelerator;
- Debugging: when you suspect a problem is related to execution timing, use it to make operations complete in a definite order;
- Synchronization points: when later non-JAX code depends on the full result (saving to disk, converting to NumPy, …).
Implicit synchronization¶
Some operations need concrete values on the host and automatically wait for the computation to finish:
- Copying to the host:
np.asarray(arr)/np.array(arr),.item(),float(arr),.tolist(); - Python control flow:
if arr > 0:,while ..., which are reallybool(arr)and need a concrete value; - Printing:
print(arr)has to format concrete values (with one important exception in 0.12, below).
Let's measure each one: dispatch a fairly heavy computation, immediately run some operation, and see how long that operation took and whether the result is ready afterward:
heavy = jax.jit(lambda m: jnp.tanh(m @ m) @ m)
heavy(x_h).block_until_ready() # warm-up
sync_ops = {
".item()": lambda r: r[0, 0].item(),
"float()": lambda r: float(r[0, 0]),
"np.asarray": lambda r: np.asarray(r),
"Python if": lambda r: bool(r[0, 0] > 0),
".tolist()": lambda r: r[:2, :2].tolist(),
}
print(f"{'operation':14s}{'dispatch(ms)':>10s}{'op(ms)':>12s}{'is_ready after':>16s}")
for name, op in sync_ops.items():
t0 = time.perf_counter(); r = heavy(x_h); t1 = time.perf_counter()
op(r); t2 = time.perf_counter()
print(f"{name:14s}{(t1-t0)*1e3:>10.3f}{(t2-t1)*1e3:>12.3f}{str(r.is_ready()):>16s}")
assert r.is_ready() # when an op that needs a host value returns, the computation must have finished
An important exception in 0.12: print/repr of large arrays doesn't synchronize. The implementation of jax.Array.__repr__ shows that for arrays of 1 MiB or more (nbytes >= 1 << 20), it prints only Array(shape=..., dtype=...) and doesn't fetch the values, so it doesn't wait for the computation to finish. Only small arrays get their values copied back to the host for printing.
r_big = heavy(x_h) # 2000×2000 float32 = 16 MB
s_big = repr(r_big)
print("large array repr:", s_big, "| is_ready after repr =", r_big.is_ready())
r_small = jax.jit(lambda m: (jnp.tanh(m @ m) @ m)[:2, :2])(x_h)
s_small = repr(r_small)
print("small array repr:", s_small.replace("\n", " "), "| is_ready after repr =", r_small.is_ready())
assert s_big == "Array(shape=(2000, 2000), dtype=float32)"
assert r_small.is_ready() and "[[" in s_small
r_big.block_until_ready()
⚠️ Note: "Printing a JAX array first copies it to the host and then formats it, so it blocks" is only true for arrays smaller than 1 MiB. For larger arrays,
reprshows only the shape and dtype, and neither fetches values nor blocks.
The course reminds us: implicit blocking exists, but using it for benchmarking is unclear and confusing. Write synchronization points explicitly with block_until_ready().
Exploiting asynchrony: overlapping host work with device computation¶
def host_work():
# Simulate host-side data preprocessing (pure NumPy, no JAX)
a = np.random.default_rng(0).random((800, 800))
return float((a @ a).sum())
t0 = time.perf_counter(); heavy(x_h).block_until_ready(); t_dev = time.perf_counter() - t0
t0 = time.perf_counter(); host_work(); t_host = time.perf_counter() - t0
t0 = time.perf_counter()
r_async = heavy(x_h) # returns right after dispatch
host_work() # the host does its own work while the device computes
r_async.block_until_ready()
t_overlap = time.perf_counter() - t0
print(f"device {t_dev*1e3:.1f} ms + host {t_host*1e3:.1f} ms = serial {(t_dev+t_host)*1e3:.1f} ms; overlapped {t_overlap*1e3:.1f} ms")
print("(in a CPU-only environment the \"device\" and the \"host\" share CPU cores, so the overlap gain is limited; it's more visible on GPU/TPU)")
🔍 Extra¶
🟢 Must-read
- Asynchronous errors show up "late": runtime errors on the device (such as OOM) may not be raised until the next synchronization, so it looks like a later line of code failed. When debugging, add
block_until_ready()at suspicious spots. - Don't call
.item()every step in a training loop. When logging the loss, synchronize every N steps, or collect the arrays andjax.device_getthem all at the end; otherwise the device sits idle waiting for the host every step. - The dispatch queue isn't infinite: dispatching too many steps at once uses more device memory to hold pending inputs and outputs. Synchronize every few steps to limit the amount of computation in flight.
- The CPU backend also dispatches asynchronously (checking
is_ready()right after dispatch above usually givesFalse); you can turn it off with the config optionjax_cpu_enable_async_dispatch(defaultTrue, applies only to non-parallel computations).
2.8 Hands-on: optimizing a numerical computation: the pairwise distance matrix¶
Computing pairwise Euclidean distances between two sets of vectors is a very common numerical task. This section puts the earlier principles to work. The goal isn't just speed, but understanding why some approaches work better in JAX/XLA.
Problem setup: $X\in\mathbb{R}^{N\times d}$ has $N$ points and $Y\in\mathbb{R}^{M\times d}$ has $M$ points. Compute the distance matrix $D\in\mathbb{R}^{N\times M}$:
$$ D_{ij}=\sqrt{\sum_{k=1}^{d}\left(X_{ik}-Y_{jk}\right)^2}. $$
(To avoid confusion with the distance matrix $D$, the dimension is written as lowercase $d$. Sizes: $N, M, d = 1000, 1500, 64$.)
This section goes through the loop "implement → time → look at jaxpr/HLO evidence → refactor" in four steps: v0 (pairwise with vmap) → v1 (the course baseline: explicit broadcasting) → v2 (the course optimization: matrix algebra) → v3 (fixing v2's numerical problem). Finally, all versions are cross-checked against float64.
# ❌ Old style (buggy code): X and Y generated with the same PRNGKey → X is exactly the first 1000 rows of Y
key_c = jax.random.PRNGKey(0)
N, M, Dd = 1000, 1500, 64
X_bad = jax.random.normal(key_c, (N, Dd))
Y_bad = jax.random.normal(key_c, (M, Dd))
same_prefix = bool(jnp.array_equal(X_bad, Y_bad[:N]))
print("X equals Y[:1000]:", same_prefix, "→ the first 1000 diagonal entries of the distance matrix are all 0; the test data is degenerate")
assert same_prefix
# ✅ Recommended: independent subkeys; put the data on the device explicitly
kx, ky = jax.random.split(jax.random.key(0))
N, M, d = 1000, 1500, 64
X = jax.device_put(jax.random.normal(kx, (N, d)))
Y = jax.device_put(jax.random.normal(ky, (M, d)))
# float64 reference answer (NumPy)
X64, Y64 = np.asarray(X, np.float64), np.asarray(Y, np.float64)
D_ref = np.sqrt(((X64[:, None, :] - Y64[None, :, :]) ** 2).sum(-1))
def report(fn, *args):
"""Compile fn and return HLO evidence + timing + error."""
c = jax.jit(fn).lower(*args).compile()
h = c.as_text()
ca, ma = c.cost_analysis(), c.memory_analysis()
return dict(t=bench(jax.jit(fn), *args), fusions=count_fusions(h), dots=h.count(" dot("),
flops=ca["flops"], bytes=ca["bytes accessed"], temp=ma.temp_size_in_bytes, hlo=h)
results = {}
v0: a pairwise function + two levels of vmap¶
The most "mathematical" version: write the distance for a single pair of points, then extend it to all pairs with two levels of vmap (for vmap, see Basics 04 Automatic vectorization with vmap).
def pair_dist(x, y):
return jnp.sqrt(jnp.sum((x - y) ** 2))
def pairwise_distance_v0(X, Y):
return jax.vmap(lambda x: jax.vmap(lambda y: pair_dist(x, y))(Y))(X)
jp0 = jax.make_jaxpr(pairwise_distance_v0)(X, Y)
print(jp0)
results["v0 vmap"] = report(pairwise_distance_v0, X, Y)
v1: the course baseline: explicit broadcasting¶
The course says that mimicking the NumPy approach in JAX uses explicit broadcasting: X[:, None, :] becomes $(N,1,d)$, Y[None, :, :] becomes $(1,M,d)$, subtracting gives a diff of shape $(N,M,d)$, and then you sum the squares along the last dimension and take the square root.
def pairwise_distance_v1(X, Y):
"""Compute pairwise distances with explicit broadcasting."""
diff = X[:, None, :] - Y[None, :, :] # X → (N,1,D), Y → (1,M,D), diff → (N,M,D)
squared_dist = jnp.sum(diff ** 2, axis=-1) # (N,M)
return jnp.sqrt(squared_dist)
jp1 = jax.make_jaxpr(pairwise_distance_v1)(X, Y)
print(jp1)
big_vars = [str(v.aval.str_short()) for e in jp1.jaxpr.eqns for v in e.outvars if v.aval.shape == (N, M, d)]
print("intermediates of shape (N,M,d) in the jaxpr:", big_vars)
results["v1 broadcast"] = report(pairwise_distance_v1, X, Y)
r0, r1 = results["v0 vmap"], results["v1 broadcast"]
print(f"v0 vs v1 after optimization: fusion {r0['fusions']} vs {r1['fusions']}; flops {r0['flops']:.3g} vs {r1['flops']:.3g}; temp {r0['temp']} vs {r1['temp']} bytes")
# Control group: force diff to be materialized with optimization_barrier and see how memory_analysis reflects it
def pairwise_distance_v1_forced(X, Y):
diff = lax.optimization_barrier(X[:, None, :] - Y[None, :, :])
return jnp.sqrt(jnp.sum(diff ** 2, axis=-1))
m_forced = jax.jit(pairwise_distance_v1_forced).lower(X, Y).compile().memory_analysis()
print(f"control: with diff forcibly materialized, temp = {m_forced.temp_size_in_bytes/1e6:.0f} MB (= N·M·d·4 bytes)")
assert len(big_vars) >= 2 # at the jaxpr level there really are (N,M,d) intermediates
assert m_forced.temp_size_in_bytes >= N * M * d * 4 # the control group shows temp can detect materialization
assert r0["flops"] == r1["flops"] and r0["fusions"] == r1["fusions"] == 1 # the vmap version and the broadcast version compile to the same thing
assert r1["temp"] < 0.01 * N * M * d * 4 # but after optimization (N,M,d) is not materialized
⚠️ Note: It's easy to assume that v1 "explicitly creates a large intermediate array
diffof shape $(N,M,D)$, with $1000\times1500\times64=96$ million elements at the example size", and to treat that as one reason v1 is slow. What actually happens (JAX 0.12, CPU): the jaxpr does contain anf32[1000,1500,64]intermediate, but XLA fuses "broadcast → subtract → square → reduce" into a single kernel.memory_analysis().temp_size_in_bytesis 0, far below the 384 MB this intermediate would need, andbytes accessedis just inputs plus output. As a control, forcing materialization withoptimization_barriergives a temp of exactly 384 MB (see theassertabove). This 384 MB array is never materialized. v1 is slow mainly because it can't use a highly optimized matmul kernel (see v2). It also has somewhat more FLOPs (about 1.4×, see the summary table), but that can't explain the several-fold timing gap in the summary table. Also, v0 (two levels of vmap) and v1 compile to the same thing;vmapis no slower than hand-written broadcasting.
v2: the course optimization: a matrix algebra identity¶
Using
$$ \lVert x_i-y_j\rVert^2=\lVert x_i\rVert^2-2\,x_i^\top y_j+\lVert y_j\rVert^2 , $$
the main cost becomes a single matmul $XY^\top$. Everything else is elementwise and can fuse with the matmul or with each other. The course uses jnp.maximum(0.0, ...) to clamp tiny negative values caused by floating-point error to 0.
def pairwise_distance_v2(X, Y):
"""Compute pairwise distances with a matrix algebra identity."""
x_sq_norms = jnp.sum(X ** 2, axis=1) # (N,)
y_sq_norms = jnp.sum(Y ** 2, axis=1) # (M,)
dot_products = jnp.dot(X, Y.T) # (N,M)
# Reshape the norms for broadcasting: (N,1) and (1,M)
squared_dist = x_sq_norms[:, None] - 2 * dot_products + y_sq_norms[None, :]
squared_dist = jnp.maximum(0.0, squared_dist) # clamp tiny negative values caused by floating-point error
return jnp.sqrt(squared_dist)
results["v2 matmul"] = report(pairwise_distance_v2, X, Y)
r2 = results["v2 matmul"]
print(clean_hlo(r2["hlo"]).split("ENTRY")[-1])
print(f"flops: v1 {r1['flops']:.4g} → v2 {r2['flops']:.4g}; dot instructions {r2['dots']}")
assert r2["dots"] == 1 and " transpose(" not in r2["hlo"] # Y.T is absorbed into dot (see 2.4)
assert r2["flops"] < r1["flops"]
v3: v2's numerical trap and the fix¶
Clamping negatives with maximum isn't enough: v2 also suffers from catastrophic cancellation. When points are far from the origin but close to each other, $\lVert x\rVert^2$, $\lVert y\rVert^2$ and $2x^\top y$ are all large numbers while their difference is small. float32 has only about 7 significant digits, so after subtraction the result is almost entirely error. v1 subtracts first and then squares, so it doesn't have this problem.
The fix is simple. Distance is translation invariant: $\lVert (x-\mu)-(y-\mu)\rVert=\lVert x-y\rVert$. Subtract a common center $\mu$ first (e.g. the mean of $X$), then compute with v2. The extra cost is only $O((N+M)d)$.
def pairwise_distance_v3(X, Y):
"""Shift to a common center first, then use the matrix identity: keeps v2's speed and fixes the cancellation error."""
mu = jnp.mean(X, axis=0)
return pairwise_distance_v2(X - mu, Y - mu)
results["v3 centered+matmul"] = report(pairwise_distance_v3, X, Y)
# Data far from the origin: +100 on every coordinate
Xo, Yo = X + 100.0, Y + 100.0
D_ref_o = np.sqrt(((np.asarray(Xo, np.float64)[:, None] - np.asarray(Yo, np.float64)[None]) ** 2).sum(-1))
err_o = {name: float(np.abs(np.asarray(jax.jit(fn)(Xo, Yo)) - D_ref_o).max())
for name, fn in [("v1", pairwise_distance_v1), ("v2", pairwise_distance_v2), ("v3", pairwise_distance_v3)]}
print("max absolute error after shifting all coordinates by +100:", {k: f"{v:.2e}" for k, v in err_o.items()})
# Self-distances: the diagonal should be 0
diag_v2 = float(np.abs(np.diag(np.asarray(jax.jit(pairwise_distance_v2)(Xo, Xo)))).max())
diag_v3 = float(np.abs(np.diag(np.asarray(jax.jit(pairwise_distance_v3)(Xo, Xo)))).max())
print(f"max of d(x, x): v2 = {diag_v2:.3g}, v3 = {diag_v3:.3g}")
assert err_o["v2"] > 1e-2 and err_o["v1"] < 1e-4 and err_o["v3"] < 1e-4
assert diag_v3 < diag_v2
Summary: evidence + timing + correctness¶
print(f"{'version':18s}{'time(ms)':>10s}{'fusion':>8s}{'dot':>5s}{'GFLOP':>9s}{'bytes(MB)':>11s}{'temp(KB)':>10s}{'max error':>11s}")
fns = {"v0 vmap": pairwise_distance_v0, "v1 broadcast": pairwise_distance_v1,
"v2 matmul": pairwise_distance_v2, "v3 centered+matmul": pairwise_distance_v3}
outs = {}
for name, r in results.items():
outs[name] = np.asarray(jax.jit(fns[name])(X, Y))
err = np.abs(outs[name] - D_ref).max()
print(f"{name:18s}{r['t']*1e3:>10.3f}{r['fusions']:>8d}{r['dots']:>5d}{r['flops']/1e9:>9.3f}"
f"{r['bytes']/1e6:>11.2f}{r['temp']/1e3:>10.1f}{err:>11.2e}")
base_t = results["v1 broadcast"]["t"]
print("\nspeedup relative to v1 (CPU timings, for reference only, no conclusions drawn):",
{k: f"{base_t / r['t']:.1f}×" for k, r in results.items()})
# Final correctness: every version matches the float64 reference, and they match each other
for name, o in outs.items():
np.testing.assert_allclose(o, D_ref, rtol=1e-5, atol=1e-4, err_msg=name)
np.testing.assert_allclose(outs["v3 centered+matmul"], outs["v1 broadcast"], rtol=1e-5, atol=1e-4)
print("✅ all four versions match the float64 reference answer")
Why v2/v3 are faster¶
Intermediate allocation: it's easy to think v1's $(N,M,d)$ intermediate eats a lot of memory bandwidth. In fact we showed above that it gets fused away and isn't materialized on CPU. On GPU, XLA usually fuses it into the reduction kernel as well (this environment is CPU-only, so that wasn't measured).- Optimized kernels: v2 relies on
jnp.dot(X, Y.T). Matmul has highly optimized implementations (YNNPACK/Eigen on CPU, cuBLAS on NVIDIA GPUs, the MXU on TPU) that get close to peak hardware throughput, whereas v1's "subtract–square–sum" is just an ordinary reduction loop. - Less work, better matched to the hardware: v2 has fewer FLOPs than v1 (see the table above), and it breaks the problem into "matmul plus elementwise ops", the combination hardware is best at.
The optimization workflow summarized by the course:
- Baseline and JIT: start from a clear, readable version (v1) and add
jax.jit; - Benchmark:
block_until_ready()plus a timer (in a notebook you can use%timeit), averaged over several runs; - Profile: capture a trace with
jax.profiler.start_trace()/stop_trace()and find the most expensive parts; use tools like Nsight Systems when you need hardware-level analysis; - Inspect the jaxpr (optional): use
jax.make_jaxprto see things before XLA optimization. This section's example shows that you also need to look at the optimized HLO andmemory_analysisto tell whether a large intermediate really exists; - Form a hypothesis and refactor: e.g. "can I use a more direct matmul instead?";
- Re-benchmark: compare the old and new versions, and verify that the results match.
The "best" implementation depends on the actual dimensions $(N,M,d)$ and the target hardware. For example, when $d$ is small (2-D or 3-D point clouds), the advantage of the matmul shrinks a lot: in rough tests on this machine's CPU, v2 is about 10× faster than v1 at $d=64$, but only 2–3× faster at $d=3$ (timings vary by machine). The v3 step also shows that optimization must not come at the expense of numerical correctness.
🔍 Extra¶
🟢 Must-read
- Be careful with $\sqrt{0}$ when differentiating distances: $\frac{d}{dz}\sqrt{z}$ is infinite at $z=0$, so
jax.gradgivesnanat self-distances (the diagonal). The usual fix is the "doublewhere" trick; see Exercise 4. For the principle, see §3.8 of Basics 03 Automatic differentiation with grad; for more approaches (includingcustom_jvp), see topic 2 in Basics 07 Practical topics. - When you only need to compare distances (kNN, cluster assignment), skip the square root: use squared distances directly. That saves the
sqrtand removes the gradient problem too. - At larger scales, compute in blocks: use
lax.map(withbatch_size) to compute row blocks and keep each block's intermediates in check. Note thatlax.mapstill assembles the full $N\times M$ result at the end; if even the output doesn't fit, reduce within each block (e.g. keep only each row's minimum distance or top-k) instead of returning the full $D$.
Hands-on exercises¶
Each of the 4 exercises below comes with a reference solution and asserts. Try writing it yourself first, then compare with the reference solution.
Exercise 1 (jaxpr detective): write a function primitive_histogram(f, *args) that returns how many times each primitive appears in the jaxpr of f (including all sub-jaxprs). Run it on target below and answer: how many times does cond appear? How many times does custom_jvp_call appear? How many sins are there in total across the sub-jaxprs?
def target(x, flag):
y = lax.cond(flag, jnp.sin, jnp.cos, x)
z = jax.jit(lambda t: jax.nn.relu(jnp.sin(t)))(y)
return jnp.sum(z)
Self-check interface: define primitive_histogram(f, *args), returning a dict or Counter of "primitive name → count". It must count recursively through all sub-jaxprs (scan bodies, while conditions and bodies, inner jit, etc.).
# ✍️ Your answer: Exercise 1
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Exercise 1
try:
primitive_histogram
except NameError as e:
print(f"⏳ Exercise 1: not answered yet: {e}")
else:
def _probe(x): # a new function: one level each of scan, while and inner jit
def body(c, _):
return jnp.sin(c) * 2.0, None
c, _ = lax.scan(body, x, None, length=3)
y, _ = lax.while_loop(lambda s: s[1] < 5, lambda s: (jnp.cos(s[0]), s[1] + 1), (c, 0))
return jax.jit(lambda t: jnp.exp(jnp.exp(t)))(y).sum()
_h = dict(primitive_histogram(_probe, jnp.ones(4)))
_expect = {"scan": 1, "sin": 1, "mul": 1, "while": 1, "lt": 1, "cos": 1,
"add": 1, "jit": 1, "exp": 2, "reduce_sum": 1}
for _name in ["sin", "cos", "lt", "exp"]:
assert _h.get(_name, 0) == _expect[_name], (
f"'{_name}' should appear {_expect[_name]} times, got {_h.get(_name, 0)}: did you miss some kind of sub-jaxpr?")
assert _h == _expect, f"counts don't match the expectation:\n got {_h}\n expected {_expect}"
print("✅ Exercise 1: self-check passed")
Exercise 2 (predicting recompilation): without running it first, predict how many times the call sequence below makes g compile, then verify with CompileCounter.
@jax.jit(static_argnames="mode")
def g(x, scale, mode="a"):
return x * scale if mode == "a" else x + scale
g(jnp.ones(3), 2.0) # ①
g(jnp.ones(3), 3.0) # ②
g(jnp.ones(3), jnp.float32(3.0)) # ③
g(jnp.ones(3), 2.0, mode="b") # ④
g(np.ones(3, np.float64), 2.0) # ⑤
g(jnp.ones((3,)), 2) # ⑥
Self-check interface: before measuring with CompileCounter, write down predicted_compiles = <your predicted number of compilations>. The self-check independently measures the same call sequence and compares it with your prediction.
# ✍️ Your answer: Exercise 2
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Exercise 2
try:
predicted_compiles
except NameError as e:
print(f"⏳ Exercise 2: not answered yet: {e}")
else:
@jax.jit(static_argnames="mode")
def _g_probe(x, scale, mode="a"): # same function as in the problem, renamed so it's counted independently
return x * scale if mode == "a" else x + scale
with CompileCounter("_g_probe") as _cc:
_g_probe(jnp.ones(3), 2.0)
_g_probe(jnp.ones(3), 3.0)
_g_probe(jnp.ones(3), jnp.float32(3.0))
_g_probe(jnp.ones(3), 2.0, mode="b")
_g_probe(np.ones(3, np.float64), 2.0)
_g_probe(jnp.ones((3,)), 2)
assert isinstance(predicted_compiles, int), "predicted_compiles should be an integer"
assert predicted_compiles == _cc.count, (
"the predicted number of compilations doesn't match the measurement. Go through the calls one by one: which properties are part of the cache key (shape, dtype, weak type, values of static arguments)?")
print("✅ Exercise 2: self-check passed")
Exercise 3 (fusion and barriers): write a function fusion_report(f, *args) that returns (number of fusions, bytes accessed). Run it on h1(t) = jnp.tanh(t * 3.0 + 1.0) ** 2 and on h2, which inserts a lax.optimization_barrier in the middle, and verify: ① h2 has more fusions than h1; ② h2's bytes accessed is larger by at least one read and one write of the intermediate array; ③ both give the same values.
Self-check interface: define fusion_report(f, *args), returning (number of fusions, bytes accessed); the count is an integer and is taken from the optimized HLO.
# ✍️ Your answer: Exercise 3
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Exercise 3
try:
fusion_report
except NameError as e:
print(f"⏳ Exercise 3: not answered yet: {e}")
else:
_h1 = lambda t: jnp.tanh(t * 3.0 + 1.0) ** 2
_h2 = lambda t: jnp.tanh(lax.optimization_barrier(t * 3.0 + 1.0)) ** 2
_t = jax.random.normal(jax.random.key(123), (256, 384))
_nbytes = _t.size * 4
(_f1, _b1), (_f2, _b2) = fusion_report(_h1, _t), fusion_report(_h2, _t)
assert isinstance(_f1, (int, np.integer)) and isinstance(_f2, (int, np.integer)), "the number of fusions should be an integer"
assert _f1 >= 1, f"h1 has {_f1} fusions: are you counting the optimized HLO?"
assert _f2 > _f1, f"there should be more fusions with the barrier: h1={_f1}, h2={_f2}"
assert _b1 >= 2 * _nbytes, f"h1 must at least read the input once and write the output once ({2 * _nbytes:,} bytes), got {_b1:,.0f}"
assert _b2 - _b1 >= 2 * _nbytes, f"h2 should add one write and one read of the intermediate array (≥ {2 * _nbytes:,} bytes), but the difference is {_b2 - _b1:,.0f}"
_f1b, _b1b = fusion_report(_h1, jnp.ones((512, 384))) # twice as many elements
assert _f1b == _f1 and 1.8 < _b1b / _b1 < 2.2, f"with twice as many elements, bytes accessed should roughly double, got {_b1:,.0f} → {_b1b:,.0f}"
print("✅ Exercise 3: self-check passed")
Exercise 4 (a differentiable safe distance): the diagonal of pairwise_distance_v1(X, X) is 0, and jax.grad gives nan there. Write a safe_pairwise_distance whose gradient is 0 where the distance is 0 and which matches v1 exactly everywhere else. Requirements: ① the forward result matches v1; ② the gradient of loss(X) = sum(D(X, X)) is finite everywhere; ③ it matches the analytic gradient in float64. The analytic gradient is $\nabla_{x_i}\sum_{a,b}\lVert x_a-x_b\rVert = 2\sum_{j\ne i}\frac{x_i-x_j}{\lVert x_i-x_j\rVert}$.
Hint: the "double where" trick first replaces 0 with a safe value before taking the square root, then uses where again to put the result back to 0.
Self-check interface: define safe_pairwise_distance(X, Y), with X of shape (N, D) and Y of shape (M, D), returning an (N, M) distance matrix.
# ✍️ Your answer: Exercise 4
# Write your code here, then run the 🧪 Self-check below. Reference solutions are in the appendix at the end of this chapter.
# 🧪 Self-check: Exercise 4
try:
safe_pairwise_distance
except NameError as e:
print(f"⏳ Exercise 4: not answered yet: {e}")
else:
_rng = np.random.default_rng(5)
_X = _rng.normal(size=(6, 4)).astype(np.float32); _X[3] = _X[1] # one pair of duplicate rows: zero distances off the diagonal too
_Y = _rng.normal(size=(5, 4)).astype(np.float32); _Y[2] = _X[0]
_D_ref = np.linalg.norm(_X.astype(np.float64)[:, None] - _Y.astype(np.float64)[None], axis=-1)
_D = safe_pairwise_distance(jnp.asarray(_X), jnp.asarray(_Y))
assert _D.shape == (6, 5), f"output shape should be (6, 5), got {_D.shape}"
assert np.allclose(_D, _D_ref, atol=1e-5), "the forward distances are wrong"
_gX = jax.jit(jax.grad(lambda Z: jnp.sum(safe_pairwise_distance(Z, Z))))(jnp.asarray(_X))
assert bool(jnp.isfinite(_gX).all()), "the gradient contains nan/inf: the gradient at zero distance isn't handled"
_Z = _X.astype(np.float64); _diff = _Z[:, None] - _Z[None]
_nrm = np.linalg.norm(_diff, axis=-1)
_unit = np.where(_nrm[..., None] > 0, _diff / np.where(_nrm > 0, _nrm, 1.0)[..., None], 0.0)
_g_ref = 2 * _unit.sum(axis=1) # zero-distance terms contribute 0
assert np.allclose(_gX, _g_ref, rtol=1e-4, atol=1e-4), "the gradient doesn't match the analytic gradient"
_gY = jax.grad(lambda Y: jnp.sum(safe_pairwise_distance(jnp.asarray(_X), Y)))(jnp.asarray(_Y))
assert bool(jnp.isfinite(_gY).all()), "nan/inf when differentiating w.r.t. Y (one row of Y equals a row of X)"
print("✅ Exercise 4: self-check passed")
Chapter Quiz¶
1. In JAX 0.12, what type does jax.make_jaxpr(f)(x) return? What are the variables left of the semicolon called, and where do they come from?
Answer and explanation
It returns a Jaxpr. 0.12 merged ClosedJaxpr into Jaxpr (jax.extend.core.ClosedJaxpr is just an alias); .jaxpr is kept as a compatibility attribute that returns the object itself, and .consts holds the actual values of the constants. In { lambda a:f32[3,3]; b:f32[3]. ... }, left of the semicolon are the constvars, which come from arrays captured by the function's closure; right of the semicolon are the invars (the function's arguments). Python scalars become Literals directly.
2. What do jax.jit(f).lower(x).as_text() and jax.jit(f).lower(x).compile().as_text() print, respectively? Which one shows fusion?
Answer and explanation
The former is StableHLO (MLIR, stablehlo.* ops), the result before XLA optimization; the latter is the optimized HLO (starting with HloModule), where you can see fusion(...) instructions, layouts like {1,0}, and backend details such as __ynn_fusion. Only the latter shows fusion.
3. Which call sequence makes jax.jit(lambda v: v * 2) compile twice? (A) f(jnp.ones(3)); f(np.ones(3, np.float64)) (B) f(1.0); f(jnp.float32(1.0)) (C) f(jnp.ones(3)); f(jnp.zeros(3))
Answer and explanation
(B). Python's 1.0 is a weak-typed float32 and jnp.float32(1.0) is a strong-typed float32, so the cache keys differ. In (A), under the default x32 mode, float64 is canonicalized to float32 and hits the cache; in (C) only the values differ.
4. In the jaxpr, v1's diff is f32[1000,1500,64] (384 MB). How can you tell whether it's actually allocated at runtime?
Answer and explanation
Look at the compiled output: jax.jit(f).lower(...).compile().memory_analysis().temp_size_in_bytes (the size of intermediate buffers), and cost_analysis()["bytes accessed"]. In practice temp is far below 384 MB and bytes is only about inputs plus output, which shows that diff was fused away and never materialized. The jaxpr alone can't answer this question.
5. Does x.T @ y inside jit materialize a transposed matrix first? How do you verify it?
Answer and explanation
No. XLA absorbs the transpose into dot's contracting dimension (lhs_contracting_dims={0}), and the optimized HLO has no transpose instruction. Data is really moved only when it must produce a new row-major array, such as returning x.T on its own.
6. What's wrong with this timing code? t0=time.perf_counter(); y=f(x); print(time.perf_counter()-t0) (f is a freshly defined jitted function)
Answer and explanation
Two problems: ① no warm-up, so the first call includes tracing and XLA compilation; ② no block_until_ready(), so asynchronous dispatch lets the function return before the computation finishes. The right way is to warm up with f(x).block_until_ready() first, then time f(x).block_until_ready(), averaged over several runs.
7. Which of these always wait for the device computation to finish? np.asarray(y), y.item(), if y > 0:, print(y) (y is 2000×2000 float32), jax.block_until_ready({"a": y})
Answer and explanation
np.asarray(y), y.item() (assuming y is a scalar; otherwise you need to index an element first), if y > 0: (assuming a scalar) and jax.block_until_ready({"a": y}) all synchronize. print(y) doesn't: in JAX 0.12, arrays of ≥ 1 MiB print only Array(shape=(2000, 2000), dtype=float32) and don't fetch values.
8. What configuration lets a script skip XLA compilation the second time it starts? Why aren't small functions written to the disk cache by default?
Answer and explanation
At the top of the script, before the first compilation, set jax.config.update("jax_compilation_cache_dir", path) (or the environment variable JAX_COMPILATION_CACHE_DIR). jax_persistent_cache_min_compile_time_secs defaults to 1.0 seconds, and functions that compile faster than that aren't written to disk, because for small functions reading from disk isn't necessarily faster than recompiling. For demos you can set it to 0. The cache key includes the jax/jaxlib versions, XLA flags and the device; if any of them changes, you get a MISS.
Version notes¶
Below is a summary of the patterns and conclusions in this chapter that matter under JAX 0.12. Each one is verified with assert at the corresponding place in the text.
| No. | Section | Key point | Where |
|---|---|---|---|
| 1 | 2.1 / 2.5 / 2.6 / 2.7 / 2.8 | Create random keys with jax.random.key(0) (typed key) plus jax.random.split, instead of the old jax.random.PRNGKey(0) | the first ❌/✅ cells in 2.1, and code in each section |
| 2 | 2.1 | The trace file isn't profile.pb but <host>.xplane.pb and <host>.trace.json.gz, optionally perfetto_trace.json.gz; TensorBoard needs the xprof plugin installed separately | 2.1 after the file-listing cell |
| 3 | 2.2 | In the jaxpr of a cond, the predicate is first convert_element_typed to int32, and the parameter is branches=(false_fun, true_fun) | 2.2 after the lax.cond cell |
| 4 | 2.2 | A Python if on a traced value inside jit raises TracerBoolConversionError right away and doesn't cause recompilation; recompilation is triggered by changes in the cache key | 2.2 "Why jaxprs matter" |
| 5 | 2.3 | The compilation pipeline is jaxpr → StableHLO (MLIR) → XLA HLO; HLO stands for High Level Operations | start of 2.3 |
| 6 | 2.3 | By default XLA doesn't simplify (x+y)-x to y (not value-preserving under IEEE); x*1.0 and x+0.0 are simplified | 2.3 algebraic simplification ❌/✅ cells |
| 7 | 2.3 | The CPU backend doesn't generate code only through LLVM: jaxlib 0.12 also uses YNNPACK library fusion (__ynn_fusion), which in this environment silently returns wrong results for "closure constant dot + reduce" | 2.3 case study |
| 8 | 2.5 | Block A (a Python if on a regular argument) raises TracerBoolConversionError on the very first call; "recompiling on the third call" never comes into it | the first ❌ cell in 2.5 |
| 9 | 2.5 | In the default x32 mode, float64 inputs are canonicalized to float32 and don't recompile; weak types (1 / 1.0 / jnp.float32(1)) do | 2.5 Strategy 2 |
| 10 | 2.5 | Creating a new jit(partial(...)) in a loop compiles every time; partial only helps if you create it once and reuse it | 2.5 Strategy 4 |
| 11 | 2.6 | Generating x and y with the same key gives two identical arrays; split first | 2.6 ❌ cell |
| 12 | 2.6 | Without jit, intermediates are jax.Arrays on the device (op-by-op dispatch), not NumPy arrays | 2.6 eager comparison |
| 13 | 2.7 | "Naive timing" without warm-up includes tracing and compilation in the first call, so it measures more than dispatch time | 2.7 ❌ timing cell |
| 14 | 2.7 | jax.block_until_ready() with no arguments raises TypeError; on a pytree it waits only for that pytree's own array leaves, not for all computation on all devices | 2.7 ❌/✅ cells |
| 15 | 2.7 | For arrays ≥ 1 MiB, print/repr shows only shape and dtype, without fetching values or blocking | 2.7 implicit synchronization |
| 16 | 2.8 | Generating X and Y with the same key gives X == Y[:1000], so the test data is degenerate | 2.8 ❌ cell |
| 17 | 2.8 | v1's 96-million-element (384 MB) intermediate exists only in the jaxpr; after XLA fusion temp_size is far below 384 MB, so it isn't materialized | 2.8 after v1 |
| 18 | 2.8 | v2's maximum(0, ·) alone isn't enough: data far from the origin suffers catastrophic cancellation (error > 1e-2); centering first fixes it (v3) | 2.8 v3 |
| 19 | 2.6 | Fusion isn't decided by XLA analyzing the jaxpr: the jaxpr is first lowered to StableHLO, and fusion happens on XLA's HLO (see the before/after_optimizations comparison in 2.3) | 2.6 paragraph on how fusion works |
Chapter summary¶
- Measure first, then optimize: capture traces with
jax.profiler.trace(warm up first,block_until_readyevery step) and view them in Perfetto or TensorBoard+XProf; label host-side time spans withTraceAnnotationand HLO op names withnamed_scope; usesave_device_memory_profileto see who's holding memory (works on CPU too). - A jaxpr shows "what JAX thinks needs computing":
invars/constvars/eqns/outvarscan all be walked programmatically; higher-order primitives (jit,cond,custom_jvp_call, …) contain nested sub-jaxprs. A large intermediate in a jaxpr doesn't mean that much memory is really allocated at runtime. - Optimized HLO shows "what XLA actually computes":
lower().as_text()is StableHLO andcompile().as_text()is optimized HLO;cost_analysis()andmemory_analysis()give FLOPs, memory traffic and intermediate buffer sizes;XLA_FLAGS=--xla_dump_todumps the output of every stage. Compilers have bugs too: when numbers look suspicious, read the HLO and cross-check against float64. - Layout: JAX is row-major by default (written
{1,0}in HLO). Inside jit,x.T @ ydoesn't materialize the transpose, but a standalonex.Tdoes. A Fortran-order NumPy array costs one conversion when it enters JAX. - Recompilation: the cache key is made of shapes, dtypes, weak types, the values of static arguments and the identity of the function object; diagnose with
jax_log_compiles/jax_explain_cache_misses; the persistent compilation cache lets new processes hit the cache directly. - Fusion is the main way to reduce memory traffic: count fusions in the HLO and quantify the gain with
bytes accessed. Reductions andoptimization_barrierform kernel boundaries. - Asynchronous dispatch: calls return immediately; timing requires a warm-up first and then
block_until_ready..item(),np.asarrayand Python control flow synchronize implicitly;printof a large array doesn't. - Optimization workflow: implement, time, read the jaxpr/HLO, form a hypothesis, refactor, time again, and verify correctness at every step. In Section 2.8, v2's speedup came with a numerical risk, which v3 fixed by centering.
Cheat sheet¶
| API | What it does |
|---|---|
jax.profiler.trace(dir, create_perfetto_trace=True) | trace as a context manager; produces xplane.pb / trace.json.gz / perfetto_trace.json.gz |
jax.profiler.start_trace(dir) / stop_trace() | start/stop a trace in function form |
jax.profiler.TraceAnnotation(name) | label a span on the host timeline |
jax.profiler.StepTraceAnnotation(name, step_num=i) | label by training step; TensorBoard aggregates per step |
jax.named_scope(name) | add a prefix to op names in the jaxpr/HLO (op_name metadata) |
jax.profiler.save_device_memory_profile(path) | save a snapshot of device memory (pprof format) |
jax.profiler.ProfileData.from_file(xplane) | parse xplane.pb in Python |
jax.make_jaxpr(f, static_argnums=...)(*args) | get a Jaxpr with consts (in 0.12 ClosedJaxpr has been merged into Jaxpr) |
cj.jaxpr.invars / constvars / eqns / outvars, cj.consts | the structural fields of a jaxpr |
jax.extend.core.jaxprs_in_params(eqn.params) | extract sub-jaxprs from higher-order primitives |
jax.extend.core.jaxpr_as_fun(cj) | execute a jaxpr as a function |
jax.jit(f).lower(*args).as_text() | StableHLO text (debug_info=True adds source locations) |
lowered.compile().as_text() | optimized HLO |
compiled.cost_analysis() | dict: flops / transcendentals / bytes accessed |
compiled.memory_analysis() | argument/output/temp_size_in_bytes, etc. |
XLA_FLAGS=--xla_dump_to=DIR --xla_dump_hlo_as_text | dump HLO at every stage (set before the process starts) |
XLA_FLAGS=--xla_disable_hlo_passes=fusion | turn off the fusion pass, for comparison |
XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type= | turn off YNNPACK library fusion on CPU (for investigating the jaxlib 0.12 bug) |
arr.format.layout | view an array's physical layout |
with jax.log_compiles(True): / jax_log_compiles | print every XLA compilation |
jax.config.update("jax_explain_cache_misses", True) | explain why the cache missed |
jax.config.update("jax_compilation_cache_dir", d) | enable the persistent compilation cache |
jax_persistent_cache_min_compile_time_secs | compile-time threshold for writing to disk (default 1 s) |
@jax.jit(static_argnums=...) | static arguments (decorator-factory form, available in 0.12) |
lax.optimization_barrier(x) | stop XLA from fusing or reordering across this point |
x.block_until_ready() / jax.block_until_ready(pytree) | wait for computation to finish (an argument is required) |
x.is_ready() | check without blocking whether the computation has finished |
jax_cpu_enable_async_dispatch | controls whether the CPU backend dispatches asynchronously |
Appendix: reference solutions¶
Write your own version in the "✍️ Your answer" cells first and check it with "🧪 Self-check"; come here to compare once the self-check passes or you're truly stuck. Running the solutions below in order overwrites any functions with the same names that you defined in the answer cells.
Exercise 1¶
For the problem, see "Hands-on exercises" in the main text.
# Reference solution: Exercise 1
def primitive_histogram(f, *args):
def walk(jaxpr, c):
for eqn in jaxpr.eqns:
c[eqn.primitive.name] += 1
for sub in je.core.jaxprs_in_params(eqn.params):
walk(sub, c)
return c
return walk(jax.make_jaxpr(f)(*args).jaxpr, Counter())
def target(x, flag):
y = lax.cond(flag, jnp.sin, jnp.cos, x)
z = jax.jit(lambda t: jax.nn.relu(jnp.sin(t)))(y)
return jnp.sum(z)
hist = primitive_histogram(target, jnp.ones(4), True)
print(dict(hist))
assert hist["cond"] == 1
assert hist["custom_jvp_call"] == 1 # relu
assert hist["sin"] == 2 # 1 in cond's true branch + 1 inside the jit
assert hist["jit"] == 2 # 1 explicit jit + 1 inside relu
Exercise 2¶
For the problem, see "Hands-on exercises" in the main text.
# Reference solution: Exercise 2
predicted_compiles = 4 # write down the prediction first, then measure with CompileCounter
@jax.jit(static_argnames="mode")
def g(x, scale, mode="a"):
return x * scale if mode == "a" else x + scale
with CompileCounter("g") as cc:
g(jnp.ones(3), 2.0) # ① new: f32[3], f32[] weak, mode=a
g(jnp.ones(3), 3.0) # ② hit: the value changed, but the type didn't
g(jnp.ones(3), jnp.float32(3.0)) # ③ new: scale is no longer weakly typed
g(jnp.ones(3), 2.0, mode="b") # ④ new: the value of the static argument changed
g(np.ones(3, np.float64), 2.0) # ⑤ hit: in x32 mode float64 is canonicalized to float32
g(jnp.ones((3,)), 2) # ⑥ new: scale becomes int32 weak
print("compilations:", cc.count)
for r in cc.records:
print(" ", r.split(" with global shapes and types ")[1].split(". Argument")[0])
assert cc.count == predicted_compiles == 4
Exercise 3¶
For the problem, see "Hands-on exercises" in the main text.
# Reference solution: Exercise 3
def fusion_report(f, *args):
c = jax.jit(f).lower(*args).compile()
return count_fusions(c.as_text()), c.cost_analysis()["bytes accessed"]
h1 = lambda t: jnp.tanh(t * 3.0 + 1.0) ** 2
h2 = lambda t: jnp.tanh(lax.optimization_barrier(t * 3.0 + 1.0)) ** 2
t_in = jax.random.normal(jax.random.key(7), (512, 512))
(f1, b1), (f2, b2) = fusion_report(h1, t_in), fusion_report(h2, t_in)
print(f"h1: fusion={f1}, bytes={b1:,.0f} | h2: fusion={f2}, bytes={b2:,.0f}")
inter = t_in.size * 4
assert f2 > f1
assert b2 - b1 >= 2 * inter # one extra write and one extra read of the intermediate array
np.testing.assert_allclose(h1(t_in), h2(t_in), rtol=1e-6)
Exercise 4¶
For the problem, see "Hands-on exercises" in the main text.
# Reference solution: Exercise 4
def safe_pairwise_distance(X, Y):
sq = jnp.sum((X[:, None, :] - Y[None, :, :]) ** 2, axis=-1)
is_zero = sq == 0.0
safe_sq = jnp.where(is_zero, 1.0, sq) # first where: keep sqrt's input away from 0
return jnp.where(is_zero, 0.0, jnp.sqrt(safe_sq)) # second where: put the result back to 0
Xs = jax.random.normal(jax.random.key(11), (40, 3))
loss_naive = lambda Z: jnp.sum(pairwise_distance_v1(Z, Z))
loss_safe = lambda Z: jnp.sum(safe_pairwise_distance(Z, Z))
g_naive = jax.jit(jax.grad(loss_naive))(Xs)
g_safe = jax.jit(jax.grad(loss_safe))(Xs)
print("naive gradient contains nan:", bool(jnp.isnan(g_naive).any()), "| safe gradient is all finite:", bool(jnp.isfinite(g_safe).all()))
# float64 analytic gradient
Z64 = np.asarray(Xs, np.float64)
diff64 = Z64[:, None, :] - Z64[None, :, :]
nrm = np.linalg.norm(diff64, axis=-1); np.fill_diagonal(nrm, np.inf)
g_ref = 2 * (diff64 / nrm[..., None]).sum(axis=1)
assert bool(jnp.isnan(g_naive).any())
assert bool(jnp.isfinite(g_safe).all())
np.testing.assert_allclose(safe_pairwise_distance(Xs, Xs), pairwise_distance_v1(Xs, Xs), atol=1e-6)
np.testing.assert_allclose(g_safe, g_ref, rtol=1e-4, atol=1e-4)
print("✅ the safe version matches in the forward pass, and its gradient matches the float64 analytic solution")