2. JAX基础数组与函数变换

JAX 的核心设计可以用一句话概括：它是一套可组合的函数变换系统，而数组库只是它的底座。理解这一点，才能明白为什么同样的代码在 JAX 里要写成“纯函数 + 变换”的形式。本节围绕三件事展开：用 jax.numpy 做数组操作、把计算写成纯函数、再把 jit 与 grad 组合起来使用。

一、jax.numpy：接口像 NumPy，语义有差异

jax.numpy（惯例导入为 jnp）的 API 与 NumPy 高度相似，jnp.array、jnp.dot、jnp.sum、切片、广播等用法基本一致。但底层数组是不可变的，任何“修改”都会返回新数组：

import jax
import jax.numpy as jnp

x = jnp.array([1.0, 2.0, 3.0])
y = x.at[1].set(10.0)      # 函数式更新，x 本身不变
print(x)                    # [1. 2. 3.]
print(y)                    # [ 1. 10.  3.]

另一个关键差异是类型与精度。JAX 默认开启 32 位模式，Python 浮点常量会被降为 float32。做梯度检查或数值敏感的计算时，需要显式开启 64 位：

jax.config.update("jax_enable_x64", True)
z = jnp.array([1.0, 2.0], dtype=jnp.float64)

此外，jnp 中的控制流（if、for、while）在变换下会被追踪，普通 Python 分支若依赖被追踪的值会报错，应改用 jnp.where、jax.lax.cond、jax.lax.scan 等表达。这是初学者最常踩的坑之一。

二、纯函数：变换的前提条件

jit 和 grad 都要求被变换的函数是纯函数：相同输入必然得到相同输出，且不产生可观察的副作用。具体来说，函数内部不应修改外部变量、不应读写文件、不应使用随机数生成器（除非显式传入 PRNGKey）、不应依赖全局可变状态。

下面是一个反例与正例的对照。反例把累加器放在函数外部，正例把状态作为参数传入并作为返回值传出：

# 反例：依赖外部可变状态
total = 0.0
def bad_add(x):
    global total
    total += x
    return total

# 正例：状态显式进出
def good_add(total, x):
    new_total = total + x
    return new_total, new_total

纯函数带来的好处是：JAX 可以安全地追踪、缓存、并行化、求导。写 JAX 代码时，应养成“所有依赖都从参数进来，所有结果都从返回值出去”的习惯。

三、jit：把函数编译成融合内核

jax.jit 把函数追踪成计算图，再交给 XLA 编译。首次调用有编译开销，之后相同形状与类型的输入会命中缓存，直接执行编译后的内核。典型用法是装饰器：

@jax.jit
def linear(w, b, x):
    return jnp.dot(x, w) + b

w = jnp.ones((3, 2))
b = jnp.zeros((2,))
x = jnp.ones((4, 3))
print(linear(w, b, x).shape)   # (4, 2)

需要注意：jit 会按输入的形状与 dtype 分别编译。如果每次调用都换形状，就会反复编译，反而更慢。实践中应尽量固定形状，或用 jax.lax 的批处理原语处理变长数据。另外，jit 内部不要打印调试信息，因为追踪阶段只会执行一次。

四、grad 与 jit 的组合：先求导再编译

jax.grad 接收一个标量输出函数，返回其梯度函数。它与 jit 可以自由嵌套，常见写法是“先 grad 再 jit”，即对梯度函数本身做编译：

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

grad_fn = jax.jit(jax.grad(loss))

w = jnp.zeros(3)
x = jnp.array([[1., 2., 3.], [4., 5., 6.]])
y = jnp.array([1., 2.])
g = grad_fn(w, x, y)
print(g)   # 梯度向量，形状 (3,)

这里 jax.grad(loss) 默认对第一个参数 w 求导。若要对其他参数求导，用 argnums 指定；若要同时得到函数值和梯度，用 jax.value_and_grad，这在训练循环中能省一次前向计算：

loss_and_grad = jax.jit(jax.value_and_grad(loss))
value, g = loss_and_grad(w, x, y)

组合顺序会影响性能与语义。jax.jit(jax.grad(f)) 表示“编译整个求导后的计算图”，通常更快；jax.grad(jax.jit(f)) 表示“对编译后的函数求导”，在 JAX 中同样可行，但一般推荐前者，因为编译器能看到完整的反向计算图并做融合优化。

五、一个可运行的最小训练片段

把上述要素串起来，就是一个线性回归的梯度下降循环。注意参数更新用 .at[].set() 或直接重新赋值，而不是原地修改：

import jax, jax.numpy as jnp

def loss(w, b, x, y):
    pred = x @ w + b
    return jnp.mean((pred - y) ** 2)

@jax.jit
def step(w, b, x, y, lr=0.1):
    val, (gw, gb) = jax.value_and_grad(loss, argnums=(0, 1))(w, b, x, y)
    return w - lr * gw, b - lr * gb, val

x = jnp.array([[1.], [2.], [3.], [4.]])
y = jnp.array([3., 5., 7., 9.])   # 目标近似 y = 2x + 1
w, b = jnp.zeros((1,)), jnp.zeros(())
for i in range(200):
    w, b, val = step(w, b, x, y)
print(w, b, val)   # 示例输出：w≈2.0, b≈1.0, loss 接近 0

判断是否达成目标，可以看两点：损失是否单调下降并趋于稳定；参数是否收敛到与数据生成规律一致的值。若损失震荡，通常是学习率过大；若损失不降，检查是否误用了非纯函数写法，或 argnums 指错了参数。

小结：jnp 提供不可变数组与函数式更新，纯函数是变换的前提，jit 负责编译加速，grad 负责自动微分，二者嵌套时优先“先求导再编译”。掌握这四点的组合方式，就具备了后续实现更复杂模型与高阶导数的基础。