入门 · 总目录

这套笔记跟着 apxml.com 的课程 JAX 入门 整理，一共 6 章，最后加了一章实战专题。每章是一个 notebook，代码都在 JAX 0.12.0.1 / Python 3.14 / CPU 上跑过，输出也保存在文件里，不运行也能直接看结果。

整套教程分两部分：入门（就是这里，第 0–7 章）和 进阶（第 0–11 章，其中第 7–11 章是补充章节）。

怎么用

  1. 在 VS Code 或 Jupyter 里打开 notebook，内核选「Python (JAX 学习)」，也就是 .venv 里的 jax-study 内核。
  2. 每章的顺序是：课程讲解 → 课程代码 → 🔍 补充 → 练习 → 测验（答案折叠着）→ 本章更新说明 → 小结与速查表 → 附录：练习参考答案。
  3. 做练习：在题目下面的「✍️ 作答区」写代码，然后运行紧跟着的「🧪 自检」。✅ 表示通过；⏳ 表示题目要的名字还没定义；别的报错会告诉你哪里不对。自检只看结果对不对，不会透露答案。参考答案都放在每章最后的附录里。

阅读约定

  • 正文代码都用 JAX 0.12 推荐的写法。已经过时的写法放在 # ❌ 旧写法 单元里，后面紧跟 # ✅ 推荐写法。
  • 需要特别留意的地方用 > ⚠️ 注意 标出来。
  • 重要结论都配了 assert。重新运行时只要不报错，就说明这些结论在你的环境里依然成立。

讲解部分是按课程要点用中文重新组织的，每节开头附有课程原文链接；测验题是根据章节内容出的。

章节索引

章Notebook核心内容本章专题更新说明
101_JAX简介JAX 与 NumPy、函数变换、数组、设备x.device() → x.device；默认 float32；异步派发；.at[] 遇到重复下标会累加24 条
202_JIT编译追踪与编译、控制流、静态参数、常见陷阱重新编译触发条件总表（8 类，每类都有编译次数断言）；分桶填充18 条
303_自动微分gradgrad、反向模式、高阶导数、value_and_gradNaN 梯度与双 where；jacfwd/jacrev 选型规则；custom_jvp18 条
404_自动向量化vmapvmap、in_axes/out_axes、嵌套、组合变换逐样本梯度（DP-SGD、找难样本、内存估算）；批处理下 cond 两个分支都会执行19 条
505_多设备并行pmapSPMD、pmap、集合通信、调试每个 pmap 示例都配一份现代写法；shard_map 梯度 N 倍专题；迁移检查清单23 条
606_状态管理纯函数、显式状态传递、PyTree、优化器状态PyTree 新 API 迁移；手写动量 SGD / Adam（与 Optax 的逐项对照见 07 专题 5）；donate_argnums21 条
707_实战专题跨章综合重新编译治理（含编译预算工具）、数值稳定、Jacobian 选型、逐样本梯度、手写优化器、内存捐赠、PyTree；综合项目：从零训练 MLP 分类器，测试准确率 99.72%，全程只编译 1 次—

环境自检

运行下面的单元，确认内核用的是 .venv 里的 JAX。

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("后端   :", jax.default_backend())
print("设备   :", jax.devices())
x = jnp.arange(3.0)
print("示例数组:", x, "| dtype:", x.dtype, "| 所在设备:", x.device)

Python : 3.14.7
JAX    : 0.12.0.1
NumPy  : 2.5.4
后端   : cpu
设备   : [CpuDevice(id=0)]
示例数组: [0. 1. 2.] | dtype: float32 | 所在设备: cpu:0

⚠️ 本环境已知问题：jaxlib 0.12.0.1 CPU 后端会悄悄算错

这是整理第 4 章时发现的。如果一个常量矩阵被 jit 函数通过闭包捕获，然后放在左边和数据相乘（W @ x）再求和，XLA:CPU 的 YNNPACK 融合会把它编译成错误的结果，而且既不报错也不警告。

  • 在本机上，n=128 时出错，n≤64 时正常。这个阈值会随形状变化（进阶第 2 章换了一组形状，n=64 就出错了）。
  • 目前只在 macOS（Apple Silicon）上复现过；在 Linux x86（Docker 模拟，接近 Colab 的环境）上测试时没有出现，相关章节在那里会自动跳过「bug 存在」的断言。
  • 写成 x @ W、或者相乘后不求和，结果都正确。
  • 不需要 vmap 也会触发：jit(lambda X: jnp.sum(W @ X, axis=0)) 就会出错，把 sum 换成 max 也一样。
  • float64 不受影响，但开启 x64 后，显式的 float32 数组照样会触发（见 进阶 11 §9.1）。
  • 这个问题从 jaxlib 0.10.2 开始出现，0.10.1 是好的（用二分法确认过）。
  • 原因是这个常量被塞进了 __ynn_fusion 内部，进阶第 2 章 对比 HLO 找到了这一点。
  • 把矩阵当作函数参数传进去，结果就是对的。
  • 在 import jax 之前设置 XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type=-reduce 也能恢复正确。

下面在子进程里对比一下（XLA 标志必须在 JAX 初始化之前设置，所以要开子进程）：

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 被闭包捕获
    as_arg  = jax.jit(jax.vmap(lambda W, x: jnp.sum(W @ x), in_axes=(None, 0)))(W, X)  # W 作为参数
    print(f"  闭包捕获 W 最大误差 = {np.abs(closure - ref).max():.3g} | W 作为参数 最大误差 = {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 最大误差 = 107 | W 作为参数 最大误差 = 5.96e-06

XLA_FLAGS='--xla_cpu_experimental_ynn_fusion_type=-reduce':
  闭包捕获 W 最大误差 = 7.63e-06 | W 作为参数 最大误差 = 7.63e-06

平时怎么避开

  • 模型参数、权重这类大数组，都当作函数参数传进去，别让 jit 通过闭包捕获。这本来就是好习惯：闭包里的数组会被当成常量写进程序，编译变慢，而且值变了也不会跟着更新。
  • 关键数值拿 NumPy float64 的结果对一下。
  • 想一劳永逸，可以在 .venv/bin/activate 里加一行 export XLA_FLAGS=--xla_cpu_experimental_ynn_fusion_type=-reduce，或者换一个 jax/jaxlib 版本再验证。

⚠️ 本环境已知问题 2：全局打开的 jax_debug_nans 在用过 vmap 后会漏检

这是整理第 7 章时发现的。用 jax.config.update("jax_debug_nans", True) 全局打开以后，只要进程里调用过一次 vmap（哪怕跟它毫无关系），之前已经跑过的 jit 函数再算出 NaN 时就不报错了，只是悄悄返回 nan。

  • 之前用过的 eager 运算也一样。比如先调用过 jnp.log，之后再对负数调用就不报错了。
  • 只有新编译的计算还能正常抓到。
  • 甚至不用真的执行 vmap，追踪过一次就够了，比如 jax.make_jaxpr(jax.vmap(...))（见 进阶 08）。
  • 用 with jax.debug_nans(True): 打开就没有这个问题，即使在同一个 with 块里先追踪过 vmap 也能正常抓到。

怎么避开（按推荐顺序）：① 用 with jax.debug_nans(True): 把要检查的那一步包起来；② 一定要全局打开的话，打开后先调用一次 jax.clear_caches()；③ 要长期留在代码里的检查，改用 jax.experimental.checkify。这三种都不受影响。详见 07_实战专题 专题 2 和 进阶 08 调试工具链。这个问题从 0.8.0 开始出现，0.7.2 是好的（用二分法确认过）。

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 写法（等价于把整段放进 with 块）
    else:
        jax.config.update("jax_debug_nans", True)     # 全局写法
    f = jax.jit(lambda x: jnp.log(x))
    f(jnp.array(2.0)).block_until_ready()            # 先正常执行一次
    if sys.argv[1] == "1":
        jax.vmap(lambda x: x * 2)(jnp.ones(3))        # 调用一次无关的 vmap
    if sys.argv[2] == "1":
        jax.clear_caches()                            # 规避方法
    try:
        print("  未报错，结果 =", f(jnp.array(-1.0)))
    except FloatingPointError:
        print("  捕获 FloatingPointError")
''')
results = {}
for vm, clr, ctx, label in [("0", "0", "0", "没用过 vmap"), ("1", "0", "0", "用过 vmap"),
                            ("1", "1", "0", "用过 vmap + clear_caches()"), ("1", "0", "1", "用过 vmap，with 写法")]:
    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 "捕获" in results["没用过 vmap"] and "未报错" in results["用过 vmap"] and "捕获" in results["用过 vmap + clear_caches()"]
assert "捕获" in results["用过 vmap，with 写法"]

没用过 vmap:
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.
  捕获 FloatingPointError

用过 vmap:
未报错，结果 = nan

用过 vmap + 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.
  捕获 FloatingPointError

用过 vmap，with 写法:
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.
  捕获 FloatingPointError

过时 API：新旧写法速查（以 JAX 0.12 为准）

旧写法在 0.12 里会怎样现在的写法章
x.device()TypeError: 'Device' object is not callablex.device（属性）、x.devices()、x.sharding1
DeviceArray / ShardedDeviceArray已不存在统一为 jax.Array1, 5
pip install jaxlib[cuda11_cudnn82]已不存在pip install "jax[cuda12]" / "jax[cuda13]" / "jax[tpu]"，要求 Python ≥ 3.121
jax.random.PRNGKey(0)能用，但属于旧接口jax.random.key(0)（类型化 key）全部
from jax.config import configModuleNotFoundErrorjax.config.update("jax_enable_x64", True)3
jax.tree_map / jax.tree_leaves …已移除jax.tree.map / jax.tree.leaves …（或 jax.tree_util.*）4, 6
jax.pmap能用、不报警告，但文档写明 “Old way…”jax.jit + NamedSharding，或 jax.shard_map5
jax.experimental.shard_mapDeprecationWarningjax.shard_map5
jax.experimental.mapsModuleNotFoundErrorjax.make_mesh + jax.sharding5
jax.device_put_replicated / device_put_shardedAttributeErrorjax.device_put(x, NamedSharding(mesh, P(...)))5
dtype=jnp.float64 就能得到 64 位被截断为 float32 并警告启动时开启 jax_enable_x641, 2

每章最容易踩的坑

  • 第 1 章：不管在什么设备上，JAX 默认都是 float32。JAX 是异步派发，不是“惰性求值”。
  • 第 2 章：函数里用 Python if 判断某个参数，就要把这个参数标成静态，只标别的参数照样报 TracerBoolConversionError。非静态参数的值变了，jit 不会重新编译，而是直接报错。
  • 第 3 章：只用 grad、不加 jit 时，函数里的 Python if 能正常工作。grad(abs)(0.0) 等于 1.0。lax.while_loop 不支持反向求导。
  • 第 4 章：默认 in_axes=0 会映射所有位置参数，不只是第一个。函数里有依赖数据的 Python if，就不能直接 vmap。对向量输出不能直接 grad。同一个 key 同时生成 uniform 和 normal，两组随机数高度相关。
  • 第 5 章：pmap 是映射轴上每个元素对应一个设备，不会自动把批量切成 N 份。把 pmean(grad) 原样搬到 shard_map 里，梯度会变成 N 倍。
  • 第 6 章：scan 的 length 必须是具体的整数，把 jit 的普通参数传给它会报 ConcretizationTypeError。把学习率标成静态参数，每次换学习率都会重新编译。

每一条的演示和正确写法，见对应章节的「🔍 补充」和章末的「本章更新说明」。

建议学习路线

  1. 按顺序学 1 → 6 章。第 2 章（jit）和第 3 章（grad）是核心；第 4 章（vmap）把前两章组合起来用；第 5 章重点看现代写法；第 6 章把前面的内容串成一个完整的训练循环。
  2. 推荐读的官方文档：🔪 JAX – The Sharp Bits，以及 Distributed arrays and automatic parallelization。
  3. 生态：Optax（优化器）→ Flax NNX（神经网络）→ Orbax（checkpoint），详见第 6 章末尾的路线图。
  4. 学完入门课后：先做 第 7 章实战专题，再进入 进阶课程。进阶第 1–6 章对应原课程（控制流、性能、分布式、进阶自动微分、自定义算子、大规模训练）；第 7–11 章是补充（Pallas 内核、调试工具链、导出与部署、随机数与分片、科学计算）。其中 进阶 08 调试工具链 和 进阶 10 随机数 建议学完入门就看。
  5. 工具：关键词索引 可以按概念或 API 查到对应章节。

💡 第一遍学习，建议按 核心学习路径 走：只学课程正文和标着「🟢 必读」的补充，标着「⚪ 选读」的可以先跳过。全部内容很多，按核心路径学能省下一大半时间。

维护：以后怎么更新这些笔记

  • 改内容：直接编辑 notebook，然后运行 .venv/bin/python tools/run_all.py -k <关键词>。只要所有 assert 都通过，说明结论在当前环境下依然成立。
  • 升级 JAX：先在新环境里跑一遍 tools/run_all.py，失败的 assert 就是行为变了的地方，逐条更新对应的说明。GitHub Actions 每周会用最新版 JAX 自动跑一次（见 .github/workflows/notebooks.yml）。
  • 重新生成索引：tools/build_index_*.py（两个目录）、tools/build_core_path.py、tools/build_keyword_index.py。改过练习或自检后再运行 tools/check_selfcheck.py。