JAX 核心特性详解：纯函数、JIT 编译、自动微分等十大必知概念-阿里云开发者社区

JAX 是 Google 和 NVIDIA 联合开发的高性能数值计算库，这两年 JAX 生态快速发展，周边工具链也日益完善了。如果你用过 NumPy 或 PyTorch，但还没接触过 JAX，这篇文章能帮助你快速上手。

围绕 JAX 已经涌现出一批好用的库：Flax 用来搭神经网络，Optax 处理梯度和优化，Equinox 提供类似 PyTorch 的接口，Haiku 则是简洁的函数式 API，Jraph 用于图神经网络，RLax 是强化学习工具库，Chex 提供测试和调试工具，Orbax 负责模型检查点和持久化。

纯函数是硬需求

JAX 对函数有个基本要求：必须是纯函数。这意味着函数不能有副作用，对同样的输入必须总是返回同样的输出。

这个约束来自函数式编程范式。JAX 内部做各种变换（编译、自动微分等）依赖纯函数的特性，用不纯的函数可能导致错误或静默失败，结果完全不对。

 # 纯函数，没问题
def pure_addition(a, b):  
  return a + b  

# 不纯的函数，JAX 不接受
counter = 0  
def impure_addition(a, b):  
  global counter  
  counter += 1  
   return a + b

JAX NumPy 与原生 NumPy

JAX 提供了类 NumPy 的接口，核心优势在于能自动高效地在 CPU、GPU 甚至 TPU 上运行，支持本地或分布式执行。这套能力来自 XLA（Accelerated Linear Algebra） 编译器，它把 JAX 代码翻译成针对不同硬件优化的机器码。

NumPy 默认只在 CPU 上跑，JAX NumPy 则不同。用法上两者很相似，这也是 JAX 容易上手的原因。


# JAX 也差不多
import jax.numpy as jnp  

print(jnp.sqrt(4))# NumPy 的写法
import numpy as np  

print(np.sqrt(4))
# JAX 也差不多
import jax.numpy as jnp  

 print(jnp.sqrt(4))

常见的操作两者看起来基本一样：

 import numpy as np  
import jax.numpy as jnp  

# 创建数组
np_a = np.array([1.0, 2.0, 3.0])  
jnp_a = jnp.array([1.0, 2.0, 3.0])  

# 元素级操作
print(np_a + 2)  
print(jnp_a + 2)  

# 广播
np_b = np.array([[1, 2, 3]])  
jnp_b = jnp.array([[1, 2, 3]])  
print(np_b + np.arange(3))  
print(jnp_b + jnp.arange(3))  

# 求和
print(np.sum(np_a))  
print(jnp.sum(jnp_a))  

# 平均值
print(np.mean(np_a))   
print(jnp.mean(jnp_a))  

# 点积
print(np.dot(np_a, np_a))   
 print(jnp.dot(jnp_a, jnp_a))

但有个重要差异需要注意：

JAX 数组是不可变的*，对数组的修改操作会返回新数组而不是改变原数组。*

NumPy 数组则可以直接修改：

 import numpy as np  

 x = np.array([1, 2, 3])  
 x[0] = 10  # 直接修改，没问题

JAX 这边就不行了：

 import jax.numpy as jnp  

 x = jnp.array([1, 2, 3])  
 x[0] = 10  # 报错

但是JAX 提供了专门的 API 来处理这种情况，通过返回一个新数组的方式实现"修改"：

 z=x.at[idx].set(y)

完整的例子：

 x = jnp.array([1, 2, 3])  
 y = x.at[0].set(10)  

 print(y)  # [10, 2, 3]  
 print(x)  # [1, 2, 3]（没变）

JIT 编译加速

即时编译（JIT）是 JAX 里一个核心特性，通过 XLA 把 Python/JAX 代码编译成优化后的机器码。

直接用 Python 解释器跑函数会很慢。加上

@jit

装饰器后，函数会被编译成快速的原生代码：

 from jax import jit  

# 不编译的版本
def square(x):  
  return x * x   

# 编译过的版本
@jit   
def jit_square(x):  
   return x * x

jit_square

快好几个数量级。函数首次调用时，JIT 引擎会：

追踪函数逻辑，构建计算图
把图编译成优化的 XLA 代码
缓存编译结果
后续调用直接用缓存的版本

自动微分

JAX 的 grad 模块能自动计算函数的导数。

 import jax.numpy as jnp  
from jax import grad  

# 定义函数：f(x) = x² + 2x + 2
def f(x):  
  return x**2 + 2 * x + 2  

# 计算导数
df_dx = grad(f)  

# 在 x = 2.0 处求值
 print(df_dx(2.0))  # 6.0

随机数处理

NumPy 用全局随机状态生成随机数。每次调用

np.random.random()

时，NumPy 会更新隐藏的内部状态：

 import numpy as np  

 np.random.random()  
 # 0.9539264374520571

JAX 的做法完全不同。作为纯函数库，它不能维护全局状态，所以要求显式传入一个伪随机数生成器（PRNG）密钥。每次生成随机数前要先分割密钥：

 from jax import random  

# 初始化密钥
key = random.PRNGKey(0)  

# 每次生成前分割
key, subkey = random.split(key)  

# 从正态分布采样
x = random.normal(subkey, ())  
print(x)  # -2.4424558  

# 从均匀分布采样
key, subkey = random.split(key)  
u = random.uniform(subkey, (), minval=0.0, maxval=1.0)  
 print(u)  # 0.104290366

一个常见的坑：同一个密钥生成的随机数始终相同。

 # 用同一个 subkey，结果重复
 x = random.normal(subkey, ())  
 print(x)  # -2.4424558  

 x = random.normal(subkey, ())  
 print(x)  # -2.4424558（还是这个值）

所以要记住总是用新密钥。

向量化：vmap

vmap 自动把函数转换成能处理批量数据的版本。逻辑上就像循环遍历每个样本，但执行效率远高于 Python 循环。

 import jax.numpy as jnp  
from jax import vmap  

def f(x):  
    return x * x + 1  

arr = jnp.array([1., 2., 3., 4.])  

# Python 循环（慢）
outputs_loop = jnp.array([f(x) for x in arr])  

# vmap 版本（快）
f_vectorized = vmap(f)  
 outputs_vmap = f_vectorized(arr)

并行化：pmap

pmap 不同于 vmap。vmap 在单个设备上做批处理，pmap 把计算分散到多个设备（GPU/TPU 核心），每个设备处理输入的一部分。

VMAP：单设备批处理向量化

PMAP：跨多设备并行执行

 import jax.numpy as jnp  
from jax import pmap  

# 查看可用设备
print(jax.devices())  # [TpuDevice(id=0), TpuDevice(id=1), ..., TpuDevice(id=7)]  

def f(x):  
    return x * x + 1  

arr = jnp.array([1., 2., 3., 4.])  

# pmap 会把数组分配到不同设备
 ys = pmap(f)(arr)

PyTrees

PyTree 在 JAX 里是个常见的概念：任何嵌套的 Python 容器（列表、字典、元组等）加上基本类型的组合。JAX 里用它来组织模型参数、优化器状态、梯度等。

 import jax.numpy as jnp  
from jax import tree_util as tu

# 构建 PyTree
pytree = {  
    "a": jnp.array([1, 2]),  
    "b": [jnp.array([3, 4]), 5]  
}  

# 获取所有叶子节点
leaves = tu.tree_leaves(pytree)  

# 对每个叶子应用函数
 doubled = tu.tree_map(lambda x: x * 2, pytree)

Optax：梯度处理和优化

Optax 是 JAX 生态里的优化库。它包含损失函数、优化器、梯度变换、学习率调度等一套工具。

损失函数：

 logits = jnp.array([[2.0, -1.0]])  
labels_onehot = jnp.array([[1.0, 0.0]])  
labels_int = jnp.array([0])  

# Softmax 交叉熵（独热编码）
loss_softmax_onehot = optax.softmax_cross_entropy(logits, labels_onehot).mean()  

# Softmax 交叉熵（整数标签）
loss_softmax_int = optax.softmax_cross_entropy_with_integer_labels(logits, labels_int).mean()  

# 二元交叉熵
loss_bce = optax.sigmoid_binary_cross_entropy(logits, labels_onehot).mean()  

# L2 损失
loss_l2 = optax.l2_loss(jnp.array([1., 2.]), jnp.array([0., 1.])).mean()  

# Huber 损失
 loss_huber = optax.huber_loss(jnp.array([1.,2.]), jnp.array([0.,1.])).mean()

优化器：

 # SGD
opt_sgd = optax.sgd(learning_rate=1e-2)  

# SGD with momentum
opt_momentum = optax.sgd(learning_rate=1e-2, momentum=0.9)   

# RMSProp
opt_rmsprop = optax.rmsprop(1e-3)  

# Adafactor
opt_adafactor = optax.adafactor(learning_rate=1e-3)  

# Adam
opt_adam = optax.adam(1e-3)  

# AdamW
 opt_adamw = optax.adamw(1e-3, weight_decay=1e-4)

梯度变换：

 # 梯度裁剪
tx_clip = optax.clip(1.0)  

# 全局梯度范数裁剪
tx_clip_global = optax.clip_by_global_norm(1.0)  

# 权重衰减（L2）
tx_weight_decay = optax.add_decayed_weights(1e-4)  

# 添加梯度噪声
 tx_noise = optax.add_noise(0.01)

学习率调度：

 # 指数衰减
 lr_exp = optax.exponential_decay(init_value=1e-3, transition_steps=1000, decay_rate=0.99)  

 # 余弦衰减
 lr_cos = optax.cosine_decay_schedule(init_value=1e-3, decay_steps=10_000)  

 # 线性预热
 lr_linear = optax.linear_schedule(init_value=0.0, end_value=1e-3, transition_steps=500)

更新步骤：

 # 计算梯度
 loss, grads = jax.value_and_grad(loss_fn)(params)  

 # 生成优化器更新
 updates, opt_state = optimizer.update(grads, opt_state)  

 # 应用更新
 params = optax.apply_updates(params, updates)

链式组合：

 # 把多个操作链起来
 optimizer = optax.chain(  
     optax.clip_by_global_norm(1.0),  # 梯度裁剪
     optax.add_decayed_weights(1e-4),  # 权重衰减
     optax.adam(1e-3)  # Adam 优化
 )

Flax 与神经网络

JAX 本身只是数值计算库，Flax 在其基础上提供了神经网络定义和训练的高级 API。Flax 代码风格接近 PyTorch，如果你用过 PyTorch 会很快上手。

Flax 提供了丰富的层和操作。基础层 包括全连接层

Dense

、卷积

Conv

、嵌入

Embed

、多头注意力

MultiHeadDotProductAttention

等：

 flax.linen.Dense(features=128)  
 flax.linen.Conv(features=64, kernel_size=(3, 3))  
 flax.linen.Embed(num_embeddings=10000, features=256)  
 flax.linen.MultiHeadDotProductAttention(num_heads=8)  
 flax.linen.SelfAttention(num_heads=8)

归一化 支持多种方式：

 flax.linen.BatchNorm()  
 flax.linen.LayerNorm()  
 flax.linen.GroupNorm(num_groups=32)  
 flax.linen.RMSNorm()

激活和 Dropout：

 flax.linen.relu(x)  
 flax.linen.gelu(x)  
 flax.linen.sigmoid(x)  
 flax.linen.tanh(x)  
 flax.linen.Dropout(rate=0.1)

池化：

 flax.linen.avg_pool(x, window_shape=(2,2), strides=(2,2))  
 flax.linen.max_pool(x, window_shape=(2,2), strides=(2,2))

循环层：

 flax.linen.LSTMCell()  
 flax.linen.GRUCell()  
 flax.linen.OptimizedLSTMCell()

下面是一个简单的多层感知机（MLP）例子：

 import jax  
import jax.numpy as jnp  
from flax import linen as nn  

class MLP(nn.Module):  
    features: list  

    @nn.compact  
    def __call__(self, x):  
        for f in self.features[:-1]:  
            x = nn.Dense(f)(x)
            x = nn.relu(x)

        x = nn.Dense(self.features[-1])(x)  
        return x  

model = MLP([32, 16, 10])  
key = jax.random.PRNGKey(0)  

# 输入：batch_size=1, 特征数=4
x = jnp.ones((1, 4))  

# 初始化参数
params = model.init(key, x)  

# 前向传播
y = model.apply(params, x)  

print("Input:", x)  
# Input: [[1. 1. 1. 1.]]  

print("Input shape:", x.shape)  
# Input shape: (1, 4)   

print("Output:", y)  
# Output: [[ 0.51415515  0.36979797  0.6212194  -0.74496573 -0.8318489   0.6590691 0.89224255  0.00737424  0.33062232  0.34577468]]  

print("Output shape:", y.shape)  
 # Output shape: (1, 10)

Flax 用

@nn.compact

装饰器，让你在

__call__

方法里直接定义层。参数是独立于模型对象存储的，需要通过

init

方法显式初始化，然后在

apply

方法中使用。

总结

JAX 的出现解决了一个长期存在的问题：如何让 Python 科学计算既保持灵活性，又能获得接近 C/CUDA 的性能。

不过 JAX 的学习曲线确实比 PyTorch 陡。纯函数的约束、不可变数组的特性、显式密钥管理等细节起初会有些别扭。但一旦习惯会发现它带来的优雅和灵活性。

https://avoid.overfit.cn/post/a16194fdc3ea450f858515d7cb3d49c4

作者：Ashish Bamania

JAX 核心特性详解：纯函数、JIT 编译、自动微分等十大必知概念

纯函数是硬需求

JAX NumPy 与原生 NumPy

JIT 编译加速

自动微分

随机数处理

向量化：vmap

并行化：pmap

PyTrees

Optax：梯度处理和优化

Flax 与神经网络

总结

热门文章

最新文章

相关电子书

探索云世界

热门

云计算

大数据

云原生

人工智能

数据库

开发与运维

活动广场

任务中心

训练营

直播

乘风者计划

下载

镜像站

技术资料

JAX 核心特性详解：纯函数、JIT 编译、自动微分等十大必知概念

纯函数是硬需求

JAX NumPy 与原生 NumPy

JIT 编译加速

自动微分

随机数处理

向量化：vmap

并行化：pmap

PyTrees

Optax：梯度处理和优化

Flax 与神经网络

总结

热门文章

最新文章

相关电子书