大规模可转换数值计算

转换 | 扩展性 | 安装指南 | 变更日志 | 参考文档

什么是 JAX？

JAX 是一个面向加速器的数组计算和程序转换的 Python 库，专为高性能数值计算和大规模机器学习设计。

JAX 可以自动对原生 Python 和 NumPy 函数进行自动微分（Automatic Differentiation）。它可以处理循环、分支、递归和闭包的微分，并且可以对导数求导。它支持通过 jax.grad 进行反向模式微分（即反向传播 Backpropagation），同时也支持前向模式微分，两者可以以任意顺序组合。

JAX 使用 XLA（一种用于加速线性代数运算的编译器）在 TPU、GPU 和其他硬件加速器上编译并扩展您的 NumPy 程序。您可以使用 jax.jit 编译您自己的纯函数。编译和自动微分也可以任意组合。

深入探究，你会发现 JAX 实际上是一个用于 可组合的函数转换 且具备 扩展性 的可扩展系统。

这是一个研究项目，并非官方 Google 产品。请预期会遇到一些 尖锐之处。请通过试用、报告错误 并提供反馈来帮助我们！

import jax
import jax.numpy as jnp

def predict(params, inputs):
  for W, b in params:
    outputs = jnp.dot(inputs, W) + b
    inputs = jnp.tanh(outputs)  # inputs to the next layer
  return outputs                # no activation on last layer

def loss(params, inputs, targets):
  preds = predict(params, inputs)
  return jnp.sum((preds - targets)**2)

grad_loss = jax.jit(jax.grad(loss))  # compiled gradient evaluation function
perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))  # fast per-example grads

目录

• 转换
• 扩展性
• 当前注意事项
• 安装
• 引用 JAX
• 参考文档

转换

核心而言，JAX 是一个用于转换数值函数的可扩展系统。这里有三个主要的转换：jax.grad、jax.jit 和 jax.vmap。

使用 grad 进行自动微分

使用 jax.grad 高效计算反向模式梯度：

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743

您可以使用 grad 进行任意阶的微分：

print(jax.grad(jax.grad(jax.grad(tanh)))(1.0))
# prints 0.62162673

您可以自由地将微分与 Python 控制流结合使用：

def abs_val(x):
  if x > 0:
    return x
  else:
    return -x

abs_val_grad = jax.grad(abs_val)
print(abs_val_grad(1.0))   # prints 1.0
print(abs_val_grad(-1.0))  # prints -1.0 (abs_val is re-evaluated)

更多详情请参见 JAX 自动微分指南 和 关于自动微分的参考文档。

使用 jit 进行编译

使用 XLA 通过 jit 将您的函数进行端到端编译，既可以作为 @jit 装饰器使用，也可以作为高阶函数使用。

import jax
import jax.numpy as jnp

def slow_f(x):
  # Element-wise ops see a large benefit from fusion
  return x * x + x * 2.0

x = jnp.ones((5000, 5000))
fast_f = jax.jit(slow_f)
%timeit -n10 -r3 fast_f(x)
%timeit -n10 -r3 slow_f(x)

使用 jax.jit 会限制函数可以使用的 Python 控制流类型；更多详情请参见关于 JIT 的控制流和逻辑运算符 的教程。

使用 vmap 进行自动向量化

vmap 沿数组轴映射一个函数。但它不仅仅是循环遍历函数应用，而是将循环下推至函数的原始操作中，例如将矩阵 - 向量乘法转换为矩阵 - 矩阵乘法以获得更好的性能。

使用 vmap 可以避免在代码中携带批量维度（batch dimensions）：

import jax
import jax.numpy as jnp

def l1_distance(x, y):
  assert x.ndim == y.ndim == 1  # only works on 1D inputs
  return jnp.sum(jnp.abs(x - y))

def pairwise_distances(dist1D, xs):
  return jax.vmap(jax.vmap(dist1D, (0, None)), (None, 0))(xs, xs)

xs = jax.random.normal(jax.random.key(0), (100, 3))
dists = pairwise_distances(l1_distance, xs)
dists.shape  # (100, 100)

通过组合 jax.vmap 与 jax.grad 和 jax.jit，我们可以获得高效的雅可比矩阵或每个样本的梯度：

per_example_grads = jax.jit(jax.vmap(jax.grad(loss), in_axes=(None, 0, 0)))

扩展性

要在数千个设备上扩展您的计算，您可以使用以下任意组合：

• 基于编译器的自动并行化，您可以像使用单个全局机器一样编程，编译器选择如何切分数据和划分计算（带有一些用户提供的约束）；
• 显式切分和自动分区，您仍然拥有全局视图，但数据切分在 JAX 类型中是显式的，可以使用 jax.typeof 进行检查；
• 手动每设备编程，您拥有数据的每设备视图和计算，并可以通过显式的集合通信进行交流。

模式	视图？	显式切分？	显式集合通信？
自动	全局	❌	❌
显式	全局	✅	❌
手动	每设备	✅	✅

from jax.sharding import set_mesh, AxisType, PartitionSpec as P
mesh = jax.make_mesh((8,), ('data',), axis_types=(AxisType.Explicit,))
set_mesh(mesh)

# parameters are sharded for FSDP:
for W, b in params:
  print(f'{jax.typeof(W)}')  # f32[512@data,512]
  print(f'{jax.typeof(b)}')  # f32[512]

shard data for batch parallelism:

inputs, targets = jax.device_put((inputs, targets), P('data'))

evaluate gradients, automatically parallelized!

gradfun = jax.jit(jax.grad(loss)) param_grads = gradfun(params, (inputs, targets))

有关更多信息，请参阅 [教程](https://docs.jax.dev/en/latest/sharded-computation.html) 和
[高级指南](https://docs.jax.dev/en/latest/advanced_guide.html)。

## 注意事项与易错点

请参阅 [注意事项笔记本](https://docs.jax.dev/en/latest/notebooks/Common_Gotchas_in_JAX.html)。

## 安装

### 支持的平台

|            | Linux x86_64 | Linux aarch64 | Mac aarch64  | Windows x86_64 | Windows WSL2 x86_64 |
|------------|--------------|---------------|--------------|----------------|---------------------|
| CPU        | 支持         | 支持          | 支持         | 支持           | 支持                |
| NVIDIA GPU | 支持         | 支持          | 不适用       | 不支持         | 实验性              |
| Google TPU | 支持         | 不适用        | 不适用       | 不适用         | 不适用              |
| AMD GPU    | 支持         | 不支持        | 不适用       | 不支持         | 实验性              |
| Apple GPU  | 不适用       | 不支持        | 实验性       | 不适用         | 不适用              |
| Intel GPU  | 实验性       | 不适用        | 不适用       | 不支持         | 不支持              |


### 安装说明

| 平台             | 说明                                                                                                    |
|------------------|-------------------------------------------------------------------------------------------------------|
| CPU              | `pip install -U jax`                                                                                  |
| NVIDIA GPU       | `pip install -U "jax[cuda13]"`                                                                        |
| Google TPU       | `pip install -U "jax[tpu]"`                                                                           |
| AMD GPU (Linux)  | 遵循 [AMD 的说明](https://github.com/jax-ml/jax/blob/main/build/rocm/README.md)。                      |
| Intel GPU        | 遵循 [Intel 的说明](https://github.com/intel/intel-extension-for-openxla/blob/main/docs/acc_jax.md)。  |

有关替代安装策略的信息，请参阅 [文档](https://docs.jax.dev/en/latest/installation.html)。
这些包括从源码编译、使用 Docker 安装、使用其他版本的 CUDA、社区支持的 conda 构建，以及一些常见问题解答。

## 引用 JAX

引用此仓库：

@software{jax2018github, author = {James Bradbury and Roy Frostig and Peter Hawkins and Matthew James Johnson and Chris Leary and Dougal Maclaurin and George Necula and Adam Paszke and Jake Vander{P}las and Skye Wanderman-{M}ilne and Qiao Zhang}, title = {{JAX}: composable transformations of {P}ython+{N}um{P}y programs}, url = {http://github.com/jax-ml/jax}, version = {0.3.13}, year = {2018}, }

在上述 bibtex 条目中，姓名按字母顺序排列，版本号旨在来自 [jax/version.py](../main/jax/version.py)，年份对应项目的开源发布。

JAX 的早期版本仅支持自动微分和编译至 XLA，描述于 [发表于 SysML 2018 的论文](https://mlsys.org/Conferences/2019/doc/2018/146.pdf) 中。我们目前正在撰写一篇更全面、更最新的论文来涵盖 JAX 的理念和能力。

## 参考文档

关于 JAX API 的详细信息，请参阅
[参考文档](https://docs.jax.dev/)。

作为 JAX 开发人员的入门指南，请参阅
[开发者文档](https://docs.jax.dev/en/latest/developer.html)。

JAX 快速上手指南

JAX 是一个用于加速器导向的数组计算和程序转换的 Python 库，专为高性能数值计算和大规模机器学习设计。它支持自动微分、XLA 编译以及多设备扩展。

1. 环境准备

系统要求

JAX 支持以下操作系统架构：

• Linux: x86_64, aarch64
• macOS: aarch64 (Apple Silicon)
• Windows: x86_64, WSL2 x86_64

硬件支持

硬件类型	支持情况
CPU	所有平台均支持
NVIDIA GPU	Linux x86_64 / aarch64 (需 CUDA)
Google TPU	Linux x86_64
AMD GPU	Linux x86_64 (需 ROCm)
Intel GPU	Linux x86_64 (实验性)

前置依赖

• Python 3.x
• pip 包管理工具

提示：在中国大陆网络环境下，建议配置国内 PyPI 镜像源以加快下载速度（例如清华源或阿里源）。

2. 安装步骤

根据目标硬件选择对应的安装命令：

CPU 版本

适用于所有平台的基础安装：

pip install -U jax

GPU 版本 (NVIDIA)

适用于 Linux 平台，需要安装对应版本的 CUDA 驱动：

pip install -U "jax[cuda13]"

TPU 版本 (Google Cloud)

适用于 TPU 环境：

pip install -U "jax[tpu]"

其他硬件

• AMD GPU (Linux): 请参考 AMD 官方说明。
• Intel GPU: 请参考 Intel 官方说明。

注意：如需更详细的安装策略（如 Docker、源码编译等），请查阅 官方文档。

3. 基本使用

JAX 的核心功能包括自动微分 (grad)、即时编译 (jit) 和向量化 (vmap)。以下是最简单的自动微分示例：

import jax
import jax.numpy as jnp

def tanh(x):
  y = jnp.exp(-2.0 * x)
  return (1.0 - y) / (1.0 + y)

grad_tanh = jax.grad(tanh)
print(grad_tanh(1.0))
# prints 0.4199743

进阶组合

您可以将 grad、jit 和 vmap 任意组合以实现高性能计算：

import jax
import jax.numpy as jnp

def loss(params, inputs, targets):
  preds = predict(params, inputs)
  return jnp.sum((preds - targets)**2)

# 编译梯度评估函数
grad_loss = jax.jit(jax.grad(loss))
# 批量计算每个样本的梯度
perex_grads = jax.jit(jax.vmap(grad_loss, in_axes=(None, 0, 0)))

更多详细 API 和使用指南，请访问 JAX 参考文档。