Basics · Overview

These notes follow the apxml.com course Getting Started with JAX. There are 6 chapters, plus one extra chapter of practical topics at the end. Each chapter is one notebook. All code was run on JAX 0.12.0.1 / Python 3.14 / CPU, and the outputs are saved in the files, so you can read the results without running anything.

The tutorial has two parts: Basics (this part, Chapters 0–7) and Advanced (Chapters 0–11, where Chapters 7–11 are extra chapters).

How to use these notes

Open a notebook in VS Code or Jupyter and pick the jax-study kernel from .venv.
Each chapter goes in this order: course explanation → course code → 🔍 Extra → exercises → quiz (answers folded) → version notes → summary and cheat sheet → appendix: reference solutions.
To do an exercise, write your code in the "✍️ Your answer" cell under the question, then run the "🧪 Self-check" cell right after it. ✅ means you passed; ⏳ means a name the exercise asks for isn't defined yet; any other error tells you what's wrong. The self-check only checks whether your result is correct and never reveals the answer. Reference solutions are in the appendix at the end of each chapter.

Conventions

The main code uses the style recommended for JAX 0.12. Outdated code goes in # ❌ Old style cells, followed right away by # ✅ Recommended.
Things that need extra attention are marked with > ⚠️ Note.
Important conclusions come with an assert. If a rerun raises no errors, those conclusions still hold in your environment.

The explanations reorganize the course's key points in my own words, and each section starts with a link to the original course page. The quiz questions are written from the chapter content.

Chapter index¶

Ch.	Notebook	Core content	Special topics	Version notes
1	01_introduction	JAX vs. NumPy, function transformations, arrays, devices	`x.device()` → `x.device`; float32 by default; asynchronous dispatch; `.at[]` accumulates on repeated indices	24 items
2	02_jit	Tracing and compilation, control flow, static arguments, common pitfalls	Table of recompilation triggers (8 kinds, each with an assertion on the compile count); bucketed padding	18 items
3	03_grad	grad, reverse mode, higher-order derivatives, value_and_grad	NaN gradients and the double `where`; how to choose jacfwd vs. jacrev; `custom_jvp`	18 items
4	04_vmap	vmap, in_axes/out_axes, nesting, composing transformations	Per-example gradients (DP-SGD, finding hard examples, memory estimates); under batching, both branches of `cond` run	19 items
5	05_pmap	SPMD, pmap, collective operations, debugging	Every pmap example comes with a modern-style version; the N× gradient in shard_map; migration checklist	23 items
6	06_state	Pure functions, explicit state passing, PyTrees, optimizer state	Migrating to the new PyTree API; hand-written momentum SGD / Adam (compared term by term with Optax in topic 5 of Chapter 07); `donate_argnums`	21 items
7	07_practical_topics	Cross-chapter review	Keeping recompilation under control (with a compile-budget tool), numerical stability, choosing a Jacobian method, per-example gradients, hand-written optimizers, buffer donation, PyTrees; capstone project: train an MLP classifier from scratch, 99.72% test accuracy, compiled only once	—

Environment check¶

Run the cell below to confirm that the kernel uses the JAX from .venv.

import sys, jax, jax.numpy as jnp, numpy as np
print("Python :", sys.version.split()[0])
print("JAX    :", jax.__version__)
print("NumPy  :", np.__version__)
print("Backend:", jax.default_backend())
print("Devices:", jax.devices())
x = jnp.arange(3.0)
print("Example array:", x, "| dtype:", x.dtype, "| device:", x.device)

Python : 3.14.7
JAX    : 0.12.0.1
NumPy  : 2.5.4
Backend: cpu
Devices: [CpuDevice(id=0)]
Example array: [0. 1. 2.] | dtype: float32 | device: cpu:0

⚠️ Known issue in this environment: the jaxlib 0.12.0.1 CPU backend silently returns wrong results¶

I found this while writing Chapter 4. If a jit function captures a constant matrix through a closure, multiplies it on the left of the data (W @ x), and then sums, the YNNPACK fusion in XLA:CPU compiles it into a wrong result, with no error and no warning.

On this machine it goes wrong at n=128 and works at n≤64. The threshold depends on the shapes (Advanced Chapter 2 uses a different set of shapes, and there it already fails at n=64).
So far it has only reproduced on macOS (Apple Silicon). In our Linux x86 test (Docker, close to a Colab environment) it did not show up, and the affected chapters automatically skip the "bug is present" assertions there.
Writing x @ W, or multiplying without summing, gives correct results.
It happens without vmap too: jit(lambda X: jnp.sum(W @ X, axis=0)) is wrong, and replacing sum with max is just as wrong.
float64 is not affected, but with x64 enabled, explicit float32 arrays still trigger it (see Advanced 11 §9.1).
The issue first appeared in jaxlib 0.10.2; 0.10.1 is fine (confirmed by bisection).
The cause is that the constant gets pulled inside __ynn_fusion. Advanced Chapter 2 found this by comparing the HLO.
Passing the matrix in as a function argument gives the correct result.
Setting XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type=-reduce before import jax also fixes it.

The comparison below runs in a subprocess (XLA flags must be set before JAX initializes, so we need a fresh process):

import subprocess, sys, os, textwrap
probe = textwrap.dedent('''
    import jax, jax.numpy as jnp, numpy as np
    W = jax.random.normal(jax.random.key(1), (128, 128)) / jnp.sqrt(128.0)
    X = jax.random.normal(jax.random.key(2), (50, 128))
    ref = (np.asarray(W, np.float64) @ np.asarray(X, np.float64).T).sum(0)
    closure = jax.jit(jax.vmap(lambda x: jnp.sum(W @ x)))(X)                          # W captured by the closure
    as_arg  = jax.jit(jax.vmap(lambda W, x: jnp.sum(W @ x), in_axes=(None, 0)))(W, X)  # W passed as an argument
    print(f"  W captured by closure: max error = {np.abs(closure - ref).max():.3g} | W as argument: max error = {np.abs(as_arg - ref).max():.3g}")
''')
for flags in ["", "--xla_cpu_experimental_ynn_fusion_type=-reduce"]:
    r = subprocess.run([sys.executable, "-c", probe], env=dict(os.environ, XLA_FLAGS=flags),
                       capture_output=True, text=True)
    print(f"XLA_FLAGS={flags!r}:\n{r.stdout.rstrip() or r.stderr[-400:]}")

XLA_FLAGS='':
  W captured by closure: max error = 107 | W as argument: max error = 5.96e-06

XLA_FLAGS='--xla_cpu_experimental_ynn_fusion_type=-reduce':
  W captured by closure: max error = 7.63e-06 | W as argument: max error = 7.63e-06

How to avoid it day to day

Pass large arrays such as model parameters and weights in as function arguments, and don't let jit capture them through a closure. That's a good habit anyway: arrays in a closure are baked into the program as constants, which slows compilation, and the program won't pick up new values when they change.
Cross-check key numbers against a NumPy float64 reference.
For a permanent fix, add export XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type=-reduce to .venv/bin/activate, or switch to another jax/jaxlib version and verify again.

⚠️ Known issue 2 in this environment: a globally enabled `jax_debug_nans` misses NaNs after `vmap` has been used¶

I found this while writing Chapter 7. After turning it on globally with jax.config.update("jax_debug_nans", True), once the process has called vmap even once (even somewhere completely unrelated), jit functions that already ran no longer raise when they produce a NaN. They just quietly return nan.

Eager operations used earlier behave the same way. For example, if you called jnp.log before, calling it on a negative number afterwards no longer raises.
Only newly compiled computations are still caught.
You don't even need to execute the vmap; tracing it once is enough, e.g. jax.make_jaxpr(jax.vmap(...)) (see Advanced 08).
Turning it on with with jax.debug_nans(True): doesn't have this problem. NaNs are caught correctly even if a vmap was traced earlier inside the same with block.

How to avoid it (in order of preference): ① wrap the step you want to check in with jax.debug_nans(True):; ② if you must turn it on globally, call jax.clear_caches() once right after; ③ for checks that stay in the code long term, use jax.experimental.checkify instead. None of these three is affected. See topic 2 of 07_practical_topics and Advanced 08: Debugging toolchain. The issue first appeared in 0.8.0; 0.7.2 is fine (confirmed by bisection).

import subprocess, sys, textwrap
probe = textwrap.dedent('''
    import sys, jax, jax.numpy as jnp
    if sys.argv[3] == "1":
        jax.debug_nans(True).__enter__()              # with style (same as putting the whole block inside a with)
    else:
        jax.config.update("jax_debug_nans", True)     # global style
    f = jax.jit(lambda x: jnp.log(x))
    f(jnp.array(2.0)).block_until_ready()            # run once normally first
    if sys.argv[1] == "1":
        jax.vmap(lambda x: x * 2)(jnp.ones(3))        # call an unrelated vmap once
    if sys.argv[2] == "1":
        jax.clear_caches()                            # workaround
    try:
        print("  no error, result =", f(jnp.array(-1.0)))
    except FloatingPointError:
        print("  caught FloatingPointError")
''')
results = {}
for vm, clr, ctx, label in [("0", "0", "0", "vmap never used"), ("1", "0", "0", "vmap used"),
                            ("1", "1", "0", "vmap used + clear_caches()"), ("1", "0", "1", "vmap used, with style")]:
    out = subprocess.run([sys.executable, "-c", probe, vm, clr, ctx], capture_output=True, text=True).stdout.strip()
    results[label] = out
    print(f"{label}:\n{out}")
assert "caught" in results["vmap never used"] and "no error" in results["vmap used"] and "caught" in results["vmap used + clear_caches()"]
assert "caught" in results["vmap used, with style"]

vmap never used:
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
  caught FloatingPointError

vmap used:
no error, result = nan

vmap used + clear_caches():
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
  caught FloatingPointError

vmap used, with style:
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
Invalid nan value encountered in the output of a jax.jit function. Calling the de-optimized version.
  caught FloatingPointError

Outdated APIs: old vs. new cheat sheet (as of JAX 0.12)¶

Old style	What happens in 0.12	Current style	Ch.
`x.device()`	`TypeError: 'Device' object is not callable`	`x.device` (attribute), `x.devices()`, `x.sharding`	1
`DeviceArray` / `ShardedDeviceArray`	No longer exists	Unified as `jax.Array`	1, 5
`pip install jaxlib[cuda11_cudnn82]`	No longer exists	`pip install "jax[cuda12]"` / `"jax[cuda13]"` / `"jax[tpu]"`, requires Python ≥ 3.12	1
`jax.random.PRNGKey(0)`	Works, but it's the legacy interface	`jax.random.key(0)` (typed key)	All
`from jax.config import config`	`ModuleNotFoundError`	`jax.config.update("jax_enable_x64", True)`	3
`jax.tree_map` / `jax.tree_leaves` …	Removed	`jax.tree.map` / `jax.tree.leaves` … (or `jax.tree_util.*`)	4, 6
`jax.pmap`	Works with no warning, but the docs call it the "Old way…"	`jax.jit` + `NamedSharding`, or `jax.shard_map`	5
`jax.experimental.shard_map`	DeprecationWarning	`jax.shard_map`	5
`jax.experimental.maps`	`ModuleNotFoundError`	`jax.make_mesh` + `jax.sharding`	5
`jax.device_put_replicated` / `device_put_sharded`	`AttributeError`	`jax.device_put(x, NamedSharding(mesh, P(...)))`	5
`dtype=jnp.float64` gives you 64-bit	Truncated to float32 with a warning	Enable `jax_enable_x64` at startup	1, 2

The most common pitfall in each chapter¶

Chapter 1: JAX defaults to float32 on every device. JAX uses asynchronous dispatch, not "lazy evaluation".
Chapter 2: if a Python if inside the function tests some argument, that argument must be marked static. Marking only other arguments still raises TracerBoolConversionError. When the value of a non-static argument changes, jit doesn't recompile; it raises an error.
Chapter 3: with plain grad and no jit, a Python if inside the function works fine. grad(abs)(0.0) is 1.0. lax.while_loop doesn't support reverse-mode differentiation.
Chapter 4: the default in_axes=0 maps all positional arguments, not just the first one. A function with a data-dependent Python if can't be vmapped directly. You can't apply grad directly to a vector output. Generating uniform and normal numbers from the same key gives two highly correlated sets of random numbers.
Chapter 5: pmap puts one element of the mapped axis on each device; it doesn't automatically split the batch into N pieces. Moving pmean(grad) into shard_map unchanged makes the gradient N times too large.
Chapter 6: the length of scan must be a concrete integer; passing it a regular jit argument raises ConcretizationTypeError. Marking the learning rate as a static argument recompiles every time you change the learning rate.

For a demo and the correct style for each item, see the "🔍 Extra" sections of the matching chapter and the "Version notes" at the end of each chapter.

Suggested learning path¶

Work through Chapters 1 → 6 in order. Chapter 2 (jit) and Chapter 3 (grad) are the core; Chapter 4 (vmap) combines the two; in Chapter 5, focus on the modern style; Chapter 6 ties everything together into a complete training loop.
Recommended official docs: 🔪 JAX – The Sharp Bits and Distributed arrays and automatic parallelization.
Ecosystem: Optax (optimizers) → Flax NNX (neural networks) → Orbax (checkpoints). See the roadmap at the end of Chapter 6.
After the Basics course: do Chapter 7: practical topics first, then move on to the Advanced course. Advanced Chapters 1–6 follow the original course (control flow, performance, distributed computing, advanced autodiff, custom operations, large-scale training); Chapters 7–11 are extras (Pallas kernels, debugging toolchain, export and deployment, random numbers and sharding, scientific computing). Of these, Advanced 08: Debugging toolchain and Advanced 10: Random numbers are worth reading right after Basics.
Tools: the keyword index lets you look up the chapter for a concept or API.

💡 On your first pass, follow the core learning path: study only the course material and the extras marked "🟢 Must-read", and skip the ones marked "⚪ Optional" for now. There's a lot of content in total, and the core path saves you well over half the time.

Maintenance: how to update these notes later¶

Changing content: edit the notebook directly, then run .venv/bin/python tools/run_all.py -k <keyword>. If all asserts pass, the conclusions still hold in the current environment.
Upgrading JAX: first run tools/run_all.py once in the new environment. Every failing assert marks a place where behavior changed; update the matching explanation one by one. GitHub Actions also runs everything weekly against the latest JAX (see .github/workflows/notebooks.yml).
Regenerating indexes: tools/build_index_*.py (one per directory), tools/build_core_path.py, tools/build_keyword_index.py. After changing exercises or self-checks, run tools/check_selfcheck.py.