一句话定义：JAX 是 Google 开源的数值计算库，语法长得像 NumPy，但它能给普通 Python 函数加上自动求导、自动编译、自动批量和自动并行这几项“外挂”。
它到底特殊在哪
JAX 的核心是函数式编程（functional programming）思路：函数没有隐藏状态，同样的输入永远给同样的输出，数组一旦创建就不可修改（immutable）。正因为函数是“干净”的，JAX 才能提供几个高阶变换（transform）：
jit：把函数整体编译成机器码，用 XLA 编译器融合算子、减少 Python 开销；grad：对函数求导，写grad(f)就得到导函数；vmap：把只处理单条样本的函数自动变成处理一批的函数，不用手写 batch 维度；- 分片/并行相关接口：把一个函数自动切到多设备、多卡上跑。
一个生活化的比方：PyTorch 像开放式厨房，边做边尝、随时改配方，灵活直观；JAX 像把菜谱写成一份标准化配方，然后交给中央厨房的统一流水线——流水线可以一次性把这道菜同时做 1000 份、同时算好“如果盐放多了会怎样”。代价是：配方必须写得很规范，中途不能偷偷往锅里加东西。
和 PyTorch 的关键差异
| 维度 | PyTorch | JAX |
|---|---|---|
| 执行方式 | 即时执行（eager），逐行运行 | 可即时，也可 jit 整体编译 |
| 状态 | 张量可变，nn.Module 持有参数 | 数组不可变，参数靠显式传入（pytree） |
| 求导 | 反向传播图按运行轨迹记录 | 对纯函数做变换式求导 |
| 随机数 | 全局随机状态 | 必须显式传 PRNG key，并手动 split |
| 生态 | 训练/部署工具链完整成熟 | 研究味更浓，Flax、Optax 等库各自成篇 |
新手最容易踩的坑是随机数：JAX 里没有“全局随机种子”，你要自己 split 出子 key 传给每个需要随机性的函数，否则循环里每次采样都一样。这不是设计失误，而是纯函数约束的必然结果——只有这样才能保证编译、并行、重放都可靠。
它在业界的真实位置
JAX 主要产自 Google Research 和 DeepMind，是 TPU 上的一等公民，也支持 GPU/CPU。它在两类场景里格外顺手：一是需要自定义求导、高阶导数、物理仿真等“数学味重”的研究；二是超大规模训练——Google 用 XLA 加自动分片（sharding）的思路，让同一份代码从小机器平滑扩到成千上万加速器，Gemini 系列模型的训练栈就大量使用 JAX（具体细节以官方技术报告为准）。AlphaFold 这类科学计算项目也是 JAX 的知名用户。
对普通从业者：如果你只是做常规业务模型，PyTorch 仍是默认选项，招聘和资料都更多；但如果你想读懂 Google 系的大模型论文、做需要精细控制并行和编译的活儿，或者想理解“函数式 + 编译器”这套思路，JAX 值得花一个周末把它的四个变换玩一遍——它改变的是你写代码的思维方式，而不只是换了个 API。