JAX — 自动微分与加速计算框架

事实

一个支持自动微分和加速编译的 Python 计算框架。

人话

JAX 像给公式叫网约车：你写数学，它自动求导，还直奔显卡高架。

用于科研训练和大实验，让 Python 跑得快。

相关概念

• Automatic Differentiation｜自动微分
• Tensor Processing Unit｜张量处理单元
• 深度学习框架 PyTorch

Autodiff
自动求导是 JAX 训练模型的关键能力。

TPU
JAX 常把计算编译后高效跑在 TPU 上。

PyTorch
JAX 和 PyTorch 都是常用深度学习框架。