2. JAX数组与NumPy互操作

JAX 的数值计算接口几乎完整复刻了 NumPy，这让熟悉 NumPy 的开发者可以低成本迁移。但 JAX 并不是 NumPy 的简单替代品：它的数组是不可变的，默认驻留在设备内存（GPU/TPU）上，并且所有运算都会被追踪以便即时编译。理解这三点的差异，是正确使用 jax.numpy 与 NumPy 互操作的前提。

jax.numpy：熟悉的接口，不同的执行语义

jax.numpy（惯例缩写为 jnp）提供了与 NumPy 高度一致的函数签名，例如 jnp.array、jnp.dot、jnp.sum、jnp.reshape 等。下面这段代码在 NumPy 和 JAX 中写法几乎相同：

import numpy as np
import jax.numpy as jnp

x_np = np.arange(6, dtype=np.float32).reshape(2, 3)
x_jax = jnp.arange(6, dtype=jnp.float32).reshape(2, 3)

print(type(x_np))   # <class 'numpy.ndarray'>
print(type(x_jax))  # <class 'jaxlib.xla_extension.ArrayImpl'>
print(x_jax.device) # 例如 cuda:0 或 cpu:0

关键区别在于：x_jax 的类型不是 numpy.ndarray，而是 JAX 自己的数组类型，它携带了设备信息和编译后的计算图。调用 jnp 函数时，JAX 会把这些操作记录为 XLA 计算图，在首次执行时编译，后续相同形状与类型的调用直接复用编译结果。因此，不要用 isinstance(x, np.ndarray) 判断 JAX 数组，应使用 isinstance(x, jnp.ndarray) 或 jax.Array。

不可变数组：没有原地修改

NumPy 允许 a[0] = 1 这样的原地赋值，而 JAX 数组一旦创建就不可更改。所有“修改”操作都会返回一个新数组：

a = jnp.array([1.0, 2.0, 3.0])
# a[0] = 9.0          # 报错：JAX arrays are immutable
b = a.at[0].set(9.0)  # 正确：返回新数组
print(a)  # [1. 2. 3.]
print(b)  # [9. 2. 3.]

.at[] 是 JAX 提供的函数式更新语法，支持 set、add、mul、min、max 等操作，例如 a.at[1].add(10.0)。这种设计使 JAX 能在编译期做更激进的内存复用与并行优化，但也意味着循环中逐元素累加不能写成原地更新，而应改用 jax.lax.scan、jnp.cumsum 或向量化表达。

设备内存管理与数据搬运

JAX 数组默认放在默认设备上（通常是 GPU）。NumPy 数组则在主机内存中。两者之间的转换会触发主机与设备之间的数据传输，这是性能敏感点。常用转换方法如下：

  • jnp.asarray(np_array)：把 NumPy 数组转为 JAX 数组，若已是 JAX 数组则原样返回，不复制。
  • np.asarray(jax_array)：把 JAX 数组转回 NumPy，会同步等待计算完成并拷贝到主机内存。
  • jax.device_put(x, device)：显式把数据放到指定设备。
  • jax.device_get(x)：显式取回主机，等价于 np.asarray 的语义。

一个常见误区是在训练循环里频繁调用 np.asarray 打印损失值，这会导致每步都强制同步并搬运数据，严重拖慢 GPU 利用率。正确做法是累积若干步后再统一取回，或使用 jax.debug.print 在编译图内部输出。

判断数据是否在设备上，可查看 x.device；判断是否已编译，可观察首次调用耗时明显高于后续调用。转换本身不报错，但性能差异会在循环中放大。

与 NumPy 的主要差异对照

维度NumPyJAX
数组可变性支持原地修改不可变，用 .at[] 返回新数组
默认位置主机内存默认设备（GPU/TPU/CPU）
执行方式即时逐条执行追踪 + XLA 编译后执行
随机数全局状态显式 PRNGKey，函数式
控制流Python 原生 if/for追踪期需用 lax.cond/lax.scan
精度默认float64默认 float32，需显式开启 x64

其中精度差异尤其容易踩坑：JAX 默认关闭 64 位，jnp.array(1.0) 是 float32。若确实需要双精度，应在导入后立即执行 jax.config.update("jax_enable_x64", True)，且该设置需在创建数组前生效。

互操作实践：在两种数组间安全切换

实际工程中，数据加载常用 NumPy 或 Python 原生结构，计算用 JAX。推荐的处理链路是：在数据管道末端用 jnp.asarray 一次性上设备，在 jit 函数内部只使用 JAX 数组，仅在需要记录或导出结果时用 np.asarray 取回。下面的 SVG 展示了这一数据流向与设备边界。

NumPy 与 JAX 数组的设备边界与转换路径 主机内存（Host） NumPy ndarray np.float32 / float64 Python 列表 / 标量 数据加载、结果导出 np.asarray(jax_arr) 取回 设备内存（GPU / TPU） jax.Array 不可变，默认 float32 jit 编译后的计算图 XLA 执行、内存复用 jnp.asarray(np_arr) 上设备 jnp.asarray / device_put 主机→设备拷贝 np.asarray / device_get 同步等待 + 设备→主机拷贝 频繁跨边界转换会强制同步，应批量处理

以一个简单的线性回归训练片段为例，演示正确的互操作方式：

import jax, jax.numpy as jnp, numpy as np

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

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

# 数据在主机侧准备，一次性上设备
x_np = np.random.randn(1024, 8).astype(np.float32)
y_np = np.random.randn(1024).astype(np.float32)
x, y = jnp.asarray(x_np), jnp.asarray(y_np)

params = (jnp.zeros((8, 1)), jnp.zeros(()))
for step in range(100):
    g = grad_fn(params, x, y)
    params = jax.tree_util.tree_map(lambda p, gp: p - 0.01 * gp, params, g)

# 训练结束后统一取回
w_host = np.asarray(params[0])
print(w_host.shape)  # (8, 1)

注意 params 的更新使用了 tree_map 生成新元组，而不是原地修改，这正是不可变语义在参数更新中的体现。整个循环中没有任何 np.asarray，避免了每步同步开销。

不要把 JAX 数组长期保存在 Python 容器中反复转换，也不要在 jit 函数内部调用 NumPy 函数——NumPy 调用无法被 XLA 追踪，会导致编译失败或静默回退到主机执行。若必须使用 NumPy 逻辑，应放在 jit 外部，或改用对应的 jnp 实现。

总结来看，JAX 与 NumPy 的互操作并不复杂，难点在于时刻意识到“数组在哪里、是否可变、是否已编译”。把转换集中在数据管道边界、用 .at[] 替代原地赋值、用 jnp 替代 np 书写计算逻辑，就能在保留 NumPy 使用习惯的同时获得 JAX 的编译加速。

← 上一章 📖 返回目录 下一章 →