Jax

一、项目简介

JAX 是 Google 开源的数值计算与机器学习框架，提供自动微分、向量化、JIT 编译与函数式变换。它以「NumPy 语法 + 编译加速」为特色，在大模型训练与科学计算领域被广泛使用，也是 Alpa 等分布式框架的底层。

二、核心特性

• 自动微分：grad / jacobian 等原生支持；
• JIT 编译：通过 XLA 加速执行；
• 向量化：vmap 自动批处理；
• 并行：pmap 支持多设备；
• 函数式 API：纯函数设计便于变换；
• 高性能，尤其适合科研。

三、适用场景

• 大模型训练与研究；
• 科学计算与数值实验；
• 需要编译加速的 ML 任务。

四、快速上手

pip install jax jaxlib
import jax.numpy as jnp
x = jnp.arange(10.0)
print(jnp.sum(x))

五、注意事项

• 函数式编程思维有一定门槛；
• 调试相对不便；
• 注意版本与硬件（GPU/TPU）匹配。