📚 课程目录第2章 JAX 数组与 NumPy 的异同

2. JAX 数组与 NumPy 的异同

如果你已经熟悉 NumPy，那么上手 jax.numpy 的第一感觉会是“几乎一模一样”：同样的 np.array、np.dot、切片、广播规则，甚至大部分函数名都直接沿用。但 JAX 并不是 NumPy 的加速版，它在数组语义、内存管理和执行模型上做了根本性的取舍。理解这些差异，是后续使用 grad、jit、vmap 的前提。

一、不可变数组：从“改一个元素”到“生成新数组”

NumPy 数组是可变对象，你可以直接对某个位置赋值：

import numpy as np
a = np.array([1.0, 2.0, 3.0])
a[0] = 10.0          # 原地修改，a 变为 [10., 2., 3.]
a += 1               # 原地加法，a 变为 [11., 3., 4.]

JAX 数组则是不可变的。下面这段代码在 JAX 中会直接报错：

import jax.numpy as jnp
b = jnp.array([1.0, 2.0, 3.0])
b[0] = 10.0          # TypeError: JAX arrays are immutable

正确做法是使用函数式更新 .at[] 语法，它返回一个新数组，原数组保持不变：

b2 = b.at[0].set(10.0)      # b2 = [10., 2., 3.]，b 仍为 [1., 2., 3.]
b3 = b.at[1].add(5.0)       # b3 = [1., 7., 3.]
b4 = b.at[0:2].set(0.0)     # 切片更新：b4 = [0., 0., 3.]

这种设计不是限制，而是为了可追踪性：每个数组都是某个计算图上的一个值，函数式更新让 JAX 能安全地对计算做变换（求导、并行、编译）。需要提醒的是，.at[] 在语义上是“复制后修改”，但在 XLA 编译后，编译器常能识别出旧数组不再被使用，从而把复制优化掉，实际开销往往接近原地修改。因此不要因为担心性能而回避 .at[]。

二、设备内存管理：数组住在哪里

NumPy 数组位于主机内存（CPU RAM），由 NumPy 自己管理。JAX 数组则默认放在加速器设备上（GPU/TPU），没有加速器时回落到 CPU。可以用 jax.devices() 查看可用设备，用 arr.devices() 查看某个数组所在设备。

设备内存是稀缺资源，因此 JAX 提供了显式的搬运接口：jax.device_put(x) 把数据送上设备，jax.device_get(x) 把结果取回主机。一个常见的性能陷阱是在训练循环里反复打印或转成 NumPy：

loss = compute_loss(params, batch)
print(loss)                 # 触发隐式同步与设备到主机的拷贝
np_loss = np.asarray(loss)  # 显式搬运，同样会打断异步流水

如果每步都这样做，GPU 的异步优势会被彻底抵消。合理做法是每隔若干步才取回一次标量做日志记录。

三、异步执行与同步：什么时候真正算完

这是 JAX 与 NumPy 差异最大、也最容易被忽视的一点。NumPy 的运算调用返回时，计算已经完成，结果是“就绪”的。JAX 的运算则是异步派发的：调用 jnp.dot(a, b) 时，JAX 只是把计算提交给设备并立即返回一个数组句柄，真正的计算可能在后台进行。

那么什么时候会强制同步？主要有三种情况：把数组转成 NumPy（np.asarray）、打印数组内容、以及使用 arr.block_until_ready() 显式等待。下面的图示梳理了从 Python 调用到设备完成的计算与同步路径：

理解异步机制对性能测量尤其重要。如果你用 time.time() 包住一段 JAX 计算，测到的可能只是“派发耗时”，而不是真实计算耗时。正确的基准写法是在计时结束前调用 block_until_ready()，并且先跑几次预热以触发编译：

import time, jax, jax.numpy as jnp

x = jnp.ones((4096, 4096))
f = jax.jit(lambda v: v @ v)

f(x).block_until_ready()          # 预热，触发编译
t0 = time.time()
y = f(x)
y.block_until_ready()             # 等待真实计算完成
print("耗时:", time.time() - t0)

四、对照速查

维度	NumPy	JAX
数组可变性	可变，支持 `a[i] = v`、`a += 1`	不可变，用 `a.at[i].set(v)` 返回新数组
默认存储位置	主机内存	加速器设备内存（无加速器则 CPU）
执行模型	同步，调用返回即完成	异步派发，需显式或隐式同步
数据类型精度	默认 float64 / int64	默认 float32 / int32，需 `jax_enable_x64` 开启 64 位
随机数	全局状态，`np.random.rand()`	显式 PRNGKey，`jax.random.normal(key, shape)`

最后补充两个容易踩的细节。其一是默认精度：JAX 为了在加速器上获得更好吞吐，默认关闭 64 位，jnp.array(1.0).dtype 是 float32；如果数值算法依赖双精度，需要在导入 JAX 前设置环境变量 JAX_ENABLE_X64=True，否则会悄悄损失精度。其二是随机数：NumPy 的全局随机状态与函数式变换不兼容，JAX 改用显式的 key，每次采样都要 split 出新 key，这一点在写训练循环时会反复遇到。

迁移 NumPy 代码到 JAX 的实用顺序：先把所有原地赋值改成 .at[].set/add，再把 np.random 换成 jax.random 加显式 key，最后检查是否有依赖 float64 的计算。三步做完，大部分纯数值代码就能直接跑在 JAX 上。

把这几条差异记牢，你就不会再把 JAX 当成“能上 GPU 的 NumPy”。不可变数组保证了变换的安全性，设备内存管理决定了数据搬运的成本，异步执行则要求你重新理解“计算何时完成”。这三者共同构成了 JAX 函数式编程模型的地基。