JAX 是 Google 推出的高性能数值计算库，主打函数式 + 可组合变换 + 硬件加速，核心是把 NumPy 易用性、自动微分、JIT 编译、GPU/TPU 并行融为一体。
JAX = NumPy 风格 + 自动微分 + XLA JIT 编译 + 多设备并行，面向机器学习 / 科学计算 / 大模型研究，DeepMind 大量项目（AlphaFold、Gemini）基于 JAX 开发。
- 类 NumPy 接口（jax.numpy）
语法和 NumPy 几乎一致，数组操作、广播、线性代数写法完全熟悉，但默认不可变数组、要求纯函数（无副作用、同输入必同输出）。
- 自动微分（grad）
一行代码算任意函数梯度 / 高阶导，支持前向 / 反向模式，是训练、优化、物理仿真的基础。
- JIT 编译（jit）
用 XLA 把 Python 函数编译成高效机器码，大幅提速，摆脱 Python 解释器 overhead。
- 自动向量化 / 并行（vmap/pmap）
vmap：自动把标量 / 单样本函数批量向量化pmap：一键多 GPU/TPU 数据并行，适合大模型 / 大规模计算
- 函数式、可组合
jit/grad/vmap 可任意组合，灵活构建复杂计算流程。
- 学术研究、强化学习、物理仿真、大模型训练 / 推理（尤其 TPU 集群）
- 追求极致性能、需要精细控制计算流程的场景
- 比 NumPy：快得多、支持 GPU/TPU、自带微分 / 并行
- 比 PyTorch/TensorFlow：更函数式、编译 / 并行更底层、研究向更强，但生态更轻、工程化工具链更少
