jax

JAX - Autograd 和 XLA（加速线性代数）

JAX 是一个框架，它将类似 NumPy 的 API 与强大的可组合函数变换系统结合起来：Grad（微分）、Jit（编译）、Vmap（向量化）和 Pmap（并行化）。

适用场景

• 需要 GPU/TPU 加速的高性能科学模拟。
• PyTorch/TF 抽象过于受限时的自定义机器学习研究。
• 为优化计算高阶导数（海森矩阵、雅可比矩阵）。
• 物理信息机器学习和可微分模拟。
• 自动向量化函数（不再需要手动分批）。
• 无需修改即可在 CPU、GPU 和 TPU 上运行相同的代码。

参考文档

官方文档：https://jax.readthedocs.io/
GitHub：https://github.com/google/jax
搜索模式：jax.numpy、jax.jit、jax.grad、jax.vmap、jax.random

核心原则

纯函数（不可变性）

JAX 基于函数式编程构建。所有函数都必须是纯函数：不得产生副作用（例如修改全局变量），并且相同输入必须返回相同输出。JAX 数组不可变。

XLA（即时编译）

JAX 使用 XLA 将 Python/NumPy 代码编译并优化为适用于特定硬件的高效机器码。

手动管理 PRNG

与 NumPy 不同，JAX 要求显式管理随机状态（密钥），以确保并行环境中的可复现性。

快速参考

安装

# CPU
pip install jax jaxlib

# GPU (Check documentation for specific CUDA versions)
pip install "jax[cuda12_pip]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html

标准导入

import jax
import jax.numpy as jnp
from jax import grad, jit, vmap, pmap, random

基本模式 - 求导并使用 JIT

import jax.numpy as jnp
from jax import grad, jit

# 1. Define a pure function
def f(x):
    return jnp.sin(x) + x**2

# 2. Transform: Create a gradient function
df_dx = grad(f)

# 3. Transform: Compile for speed
f_fast = jit(f)

# 4. Use
val = f_fast(2.0)
slope = df_dx(2.0)

关键规则

✅ 应该做

• 使用 jax.numpy (jnp) - 它与 NumPy 类似，但支持 JAX 变换。
• 编写纯函数 - 确保函数只依赖输入，且不修改外部状态。
• 手动处理 PRNG 密钥 - 每次随机操作都使用 key, subkey = random.split(key)。
• 使用 vmap 进行分批处理 - 编写处理单个样本的代码，并让 JAX 处理批次维度。
• 在 JIT 中设置 static_argnums - 如果 JIT 函数接受非数组参数（例如循环中使用的字符串或整数），请将其标记为静态参数。
• 通过 .at 进行原地更新 - 由于数组不可变，请使用 x = x.at[idx].set(val)。

❌ 不应该做

• 使用原地更新 - x[idx] = val 会引发错误。
• 使用标准 numpy (np) - 标准 NumPy 数组不支持 JAX 变换。
• 使用副作用 - 不要在 JIT 编译的函数中使用 print() 或修改全局变量。
• 忘记调用 block_until_ready() - JAX 使用异步执行。要准确计时，请使用 result.block_until_ready()。

反模式（绝对不要）

import jax.numpy as jnp
from jax import jit, random

# ❌ BAD: Modifying a global variable inside a function
counter = 0
@jit
def bad_func(x):
    global counter
    counter += 1 # ❌ Side effect! Will only run once during compilation
    return x * 2

# ❌ BAD: Standard NumPy random (not reproducible/parallel-safe)
# val = np.random.randn(5) 

# ✅ GOOD: JAX PRNG
key = random.key(42)
val = random.normal(key, (5,))

# ❌ BAD: In-place assignment
# x[0] = 1.0 

# ✅ GOOD: Functional update
x = jnp.zeros(5)
x = x.at[0].set(1.0)

函数变换

Grad（微分）

def loss(params, x, y):
    pred = jnp.dot(x, params)
    return jnp.mean((pred - y)**2)

# Gradient of loss with respect to the 1st argument (params)
grads = grad(loss)(params, x, y)

# Higher-order: Hessian
hessian = jax.hessian(loss)(params, x, y)

Jit（即时编译）

@jit
def complex_math(x):
    # This whole block is compiled into one XLA kernel
    y = jnp.exp(x)
    return jnp.sin(y) / jnp.sqrt(x)

# First call: Compiles (slow)
# Subsequent calls: Super fast

Vmap（自动向量化）

def model(params, x):
    return jnp.dot(params, x)

# model works on 1D x. How to apply to a 2D batch of X?
# in_axes=(None, 0): don't map params, map the 0th axis of x
batch_model = vmap(model, in_axes=(None, 0))

batch_preds = batch_model(params, X_batch)

随机数（jax.random）

状态管理

key = random.key(0)

# Never reuse the same key!
key, subkey = random.split(key)
data = random.normal(subkey, (10,))

key, subkey = random.split(key)
noise = random.uniform(subkey, (10,))

使用 PyTrees

处理复杂数据结构（字典、列表、元组）

JAX 变换适用于“PyTrees”——由数组构成的嵌套容器。

params = {'weights': jnp.ones((5,)), 'bias': 0.0}

def predict(p, x):
    return jnp.dot(x, p['weights']) + p['bias']

# grad and jit handle the dictionary automatically
grads = grad(predict)(params, x)

实用工作流

1. 可微分物理：求解简单 ODE

def system_dynamics(state, t):
    # Simple harmonic oscillator
    x, v = state
    dxdt = v
    dvdt = -0.5 * x
    return jnp.array([dxdt, dvdt])

def loss_fn(initial_state, target_x):
    # Simulate for 10 steps using simple Euler
    state = initial_state
    dt = 0.1
    for i in range(10):
        state += system_dynamics(state, i*dt) * dt
    return (state[0] - target_x)**2

# We can take the gradient of the simulation with respect to initial state!
optimize_initial_state = grad(loss_fn)

2. 使用 vmap 进行参数扫描

def simulation(param):
    # Some complex computation
    return jnp.sum(jnp.linspace(0, param, 100))

# Parallelize simulation over a range of parameters
params = jnp.linspace(1, 10, 100)
results = vmap(simulation)(params)

3. 使用 pmap 进行分布式训练

# pmap replicates the function across multiple GPUs
# (assuming 8 GPUs are available)
# x = jnp.zeros((8, 1024))
# results = pmap(jnp.sin)(x)

性能优化

JIT 中的静态参数

如果函数使用基于输入值的循环，则该值必须是静态的。

from functools import partial

@partial(jit, static_argnums=(1,))
def power_loop(x, n):
    for i in range(n):
        x = x * x
    return x

避免 Python 控制流

优先使用 JAX 控制流（cond、while_loop、fori_loop），以便更好地进行 XLA 优化。

from jax.lax import cond

def safe_divide(x, y):
    return cond(y == 0, lambda _: 0.0, lambda _: x / y, operand=None)

常见问题及解决方案

“Tracer”错误

在 JIT 内部，JAX 看不到实际数字，而是看到“Tracers”。

# ❌ Problem:
# @jit
# def func(x):
#     if x > 0: return x # Error! JAX doesn't know x's value during compile

# ✅ Solution:
# Use jax.lax.cond or mark x as static_argnum

NaN 梯度

如果函数存在奇点（例如 sqrt(0)），梯度将为 NaN。

# ✅ Solution: Add a small epsilon
def safe_sqrt(x):
    return jnp.sqrt(x + 1e-8)

GPU 上的内存泄漏

JAX 默认会预分配 90% 的 GPU 内存。

# ✅ Solution: Set environment variable
import os
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'

最佳实践

1. 始终使用纯函数 - 不产生副作用，输出具有确定性
2. 显式管理 PRNG 密钥 - 为每个随机操作拆分密钥
3. 对热点循环使用 JIT - 编译重复调用的函数
4. 利用 vmap 进行批处理 - 编写单样本代码，让 JAX 处理批次
5. 标记静态参数 - 在 JIT 中对非数组参数使用 static_argnums
6. 使用函数式更新 - 始终使用 .at 方法修改数组
7. 先分析性能，再优化 - 使用 jax.profiler 查找瓶颈
8. 处理设备放置 - 使用 jax.device_put() 控制数组所在位置
9. 先在 CPU 上测试 - 在 CPU 上调试，然后扩展到 GPU/TPU
10. 了解编译成本 - 首次调用 JIT 较慢，后续调用很快

JAX 是可微科学的终极试验场。通过将数学视为一系列函数式变换，它释放了此前在 Python 中无法实现的速度和复杂性。