一、JAX 初探：它到底是什么？

说实话，我第一次接触 JAX 的时候，第一反应是——这不就是 NumPy 换了个马甲吗？

后来踩了几个坑，才明白事情没那么简单。JAX 本质上是一个 可微分编程框架，它把 NumPy 的接口拿过来，但底层完全重写了。你可以把它理解成「会算梯度的 NumPy」。

我个人习惯把 JAX 看作三个东西的合体：

• NumPy 的替代品——API 几乎一模一样
• 自动微分引擎——能算任意函数的梯度
• 转换器集合——jit、vmap、pmap 这些黑科技

嗯，这里要注意：JAX 不是深度学习框架，至少不完全是。它更像一个底层工具，PyTorch 和 TensorFlow 能做的事它都能做，但方式完全不同。

它和 NumPy 到底有什么关系？

直接看代码最清楚：

import jax.numpy as jnp
import numpy as np

# 创建数组的语法完全一样
a_np = np.array([1.0, 2.0, 3.0])
a_jax = jnp.array([1.0, 2.0, 3.0])

# 运算也几乎一样
print(np.sum(a_np))      # 6.0
print(jnp.sum(a_jax))    # 6.0

# 但有个关键区别
print(type(a_np))   # <class 'numpy.ndarray'>
print(type(a_jax))  # <class 'jaxlib.xla_extension.ArrayImpl'>

看出来了吗？JAX 的数组不是 NumPy 数组，而是 XLA 编译器的内部表示。这意味着什么？

我在项目中遇到过一个问题：用 NumPy 写了一个前向传播，跑得挺欢。换成 JAX 后，同样的代码，速度反而慢了。后来才发现，JAX 的数组默认在 GPU 上，而我当时根本没装 CUDA……

核心设计哲学：三个关键词

JAX 的设计哲学，说白了就三个词：函数式、不可变性、转换。我一个个讲。

1. 函数式编程

JAX 要求你把计算写成 纯函数——输入决定输出，没有副作用。什么意思？

# 不推荐：有副作用的写法
def bad_update(params, lr, grads):
    params -= lr * grads  # 修改了外部变量
    return params

# 推荐：纯函数写法
def good_update(params, lr, grads):
    new_params = params - lr * grads  # 创建新变量
    return new_params

你想想看，为什么 JAX 要这么设计？因为只有纯函数才能被 任意转换。jit 编译、自动微分、向量化映射……这些转换都假设函数没有隐藏状态。

我记得刚开始写 JAX 代码时，总忍不住写 global 变量，结果 grad 函数死活算不对。后来养成了习惯：所有状态都显式传参。

2. 不可变性

JAX 的数组一旦创建就不能修改。这和 NumPy 完全不同：

# NumPy：可以原地修改
arr = np.array([1, 2, 3])
arr[0] = 99  # 没问题

# JAX：会报错
arr = jnp.array([1, 2, 3])
arr[0] = 99  # TypeError: JAX arrays are immutable

为什么会这样？因为不可变性是函数式编程的基石。如果数组可以原地修改，JAX 就无法保证 计算图的可追踪性。

实际项目中，我经常用 jnp.where 或 jnp.concatenate 来「模拟」修改操作。虽然看起来多了一步，但换来了自动微分的确定性。

3. 转换（Transformations）

这是 JAX 最酷的部分。它提供了四个核心转换器：

转换器	作用	类比
jit	即时编译，加速计算	把 Python 函数编译成 XLA 指令
grad	自动微分，计算梯度	PyTorch 的 backward()
vmap	自动向量化，批量处理	手动写 for 循环的替代
pmap	并行计算，多设备	分布式训练的简化版

这些转换器可以 任意组合。比如：

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

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

# 先求梯度，再编译，再向量化
grad_fn = grad(loss_fn)           # 自动微分
compiled_grad = jit(grad_fn)      # 编译加速
batched_grad = vmap(compiled_grad, in_axes=(None, 0, 0))  # 批量处理

你想想看，这三行代码在 PyTorch 里要写多少？手动写 for 循环、手动管理设备、手动处理 batch……JAX 用转换器一步到位。

知识体系总览

下面这张图是我自己整理的 JAX 核心逻辑，你看一眼就能明白整体架构：

这张图我画了好几次才满意。你看，JAX 的三大支柱是相互支撑的：函数式编程保证了转换器的可靠性，不可变性保证了计算图的可追踪性，而转换器则是 JAX 的真正价值所在。

一个完整的例子

最后，给你看一个完整的 JAX 工作流。这是我做线性回归时的标准模板：

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

# 1. 初始化参数
key = random.PRNGKey(42)
params = random.normal(key, (2,))

# 2. 定义纯函数
def predict(params, x):
    return params[0] + params[1] * x

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

# 3. 自动微分 + 编译
grad_fn = jax.grad(loss)
compiled_grad = jax.jit(grad_fn)

# 4. 训练循环（纯函数式）
def update(params, x, y, lr=0.01):
    grads = compiled_grad(params, x, y)
    return params - lr * grads

# 5. 模拟数据
x = jnp.linspace(0, 10, 100)
y = 3.0 + 2.0 * x + random.normal(key, (100,)) * 0.1

# 6. 训练
for epoch in range(100):
    params = update(params, x, y)
    if epoch % 20 == 0:
        current_loss = loss(params, x, y)
        print(f"Epoch {epoch}: loss = {current_loss:.4f}")

注意看，整个过程中没有 .backward()，没有 optimizer.step()，没有 with torch.no_grad()。就是纯函数 + 转换器，干净利落。

---

公众号:蓝海资料掘金营，微信deep3321