第3章 JAX核心

3、JAX核心：jnp.array、设备内存管理、即时编译（jit）入门

好，咱们正式开始接触JAX的核心机制了。说实话，很多初学者一上来就被JAX的“函数式”和“设备管理”搞懵了。我个人觉得，这其实是因为没搞懂三个最基础的东西：jnp.array、设备内存、以及jit。今天咱们就把这三个东西彻底讲透。

3.1 jnp.array：不只是NumPy的“复制品”

你可能会想：“不就是个数组吗？跟NumPy有什么区别？”嗯，区别大了去了。

jnp.array 是JAX版本的数组。它长得像NumPy，用起来也像NumPy，但骨子里完全不同。我刚开始用的时候，也以为它就是NumPy的平替，结果踩了不少坑。

核心区别：jnp.array 是不可变的（immutable）。你没法原地修改它。比如 x[0] = 1 这种操作，在NumPy里没问题，但在JAX里会报错。

为什么会这样？因为JAX的设计哲学是“函数式”。说白了，就是每个操作都返回一个新数组，不改变原来的。这跟分布式计算里的“无副作用”要求是吻合的。

import jax.numpy as jnp
import numpy as np

# NumPy 风格（可变）
x_np = np.array([1, 2, 3])
x_np[0] = 100  # 没问题

# JAX 风格（不可变）
x_jax = jnp.array([1, 2, 3])
# x_jax[0] = 100  # 会报错！TypeError: '' does not support item assignment

# 正确做法：用索引更新函数
x_jax_new = x_jax.at[0].set(100)
print(x_jax_new)  # [100, 2, 3]

我在项目中遇到过这种情况：有个同事把NumPy代码直接搬过来，结果跑分布式训练时，因为数组不可变的问题，调试了一整天。所以，记住：用 .at[] 方法来更新元素。

3.2 设备内存管理：数据到底在哪？

JAX的一个核心能力是“设备无关”。你的数据可以在CPU上，也可以在GPU或TPU上。但这也带来了一个问题：你得知道数据在哪。

我个人习惯用 jax.devices() 先看看当前环境有哪些设备。

import jax

print(jax.devices())  
# 输出示例：[TFRT_CPU_0, TFRT_GPU_0]  （如果有GPU的话）

然后，你可以用 jax.device_put() 把数据显式放到某个设备上。

import jax.numpy as jnp

# 创建一个数组，默认在CPU上
x = jnp.array([1, 2, 3])

# 显式放到GPU上（如果有）
gpu_device = jax.devices('gpu')[0]
x_gpu = jax.device_put(x, gpu_device)

print(x_gpu.device())  # 输出：TFRT_GPU_0

注意：如果你把数据放在GPU上，但后续操作在CPU上执行，JAX会自动帮你做数据传输。但这会带来额外的开销。我曾经因为没注意这个，导致训练速度反而变慢了。所以，尽量让数据一直待在同一个设备上。

还有一个常见的坑：内存泄漏。JAX的数组是惰性求值的，也就是说，你创建了一个数组，它可能还没真正分配内存。直到你真正用到它时，内存才会被分配。这听起来很美好，但如果你不小心创建了大量中间数组，内存可能会爆掉。

我曾经在循环里反复创建 jnp.array，结果内存直接飙到几十GB。后来发现，用 jax.lax.scan 或者 jax.vmap 来替代循环，能大大减少中间数组的创建。

3.3 即时编译（jit）入门：让代码飞起来

JAX最吸引人的地方之一就是 jit。它能把你的Python函数编译成高效的XLA（加速线性代数）代码，然后在GPU或TPU上运行。说白了，就是让你的代码跑得飞快。

用法很简单：用 @jax.jit 装饰器。

import jax
import jax.numpy as jnp

@jax.jit
def f(x, y):
    return jnp.dot(x, y) + jnp.sum(x)

x = jnp.array([1.0, 2.0, 3.0])
y = jnp.array([4.0, 5.0, 6.0])

result = f(x, y)
print(result)  # 输出：37.0

第一次调用 f 时，JAX会编译它。这个过程会慢一点。但后续调用就快多了，因为编译结果被缓存了。

小技巧：如果你想知道编译过程发生了什么，可以用 jax.jit(f).lower(x, y).compile() 来手动触发编译，并查看编译后的HLO（高级优化器）代码。这在调试性能问题时很有用。

但 jit 也有它的限制。比如，它要求函数是“纯函数”——没有副作用，不能修改全局变量，不能打印东西。我刚开始用的时候，在 jit 函数里加了个 print 语句，结果发现它只在编译时打印一次，运行时根本不打印。嗯，这就是副作用被“优化”掉了。

还有一个常见的坑：动态形状。JAX的 jit 要求输入的形状是固定的。如果你传入不同形状的数组，JAX会重新编译一次。这会导致性能下降。所以，尽量保持输入形状一致。

@jax.jit
def g(x):
    return jnp.sum(x)

# 第一次调用，形状 (3,)
g(jnp.array([1, 2, 3]))  # 编译

# 第二次调用，形状 (4,)
g(jnp.array([1, 2, 3, 4]))  # 重新编译！慢！

我在项目中遇到过这种情况：有个模型训练时，每次batch size不一样，结果每次迭代都要重新编译，训练速度慢得离谱。后来我们固定了batch size，并用padding补齐，问题就解决了。

3.4 知识体系图：三者如何协作？

下面这张图展示了 jnp.array、设备内存管理、jit 三者之间的关系。你可以看到，数据从创建到计算，再到结果输出，每一步都涉及设备管理和编译优化。

3.5 避坑指南：我踩过的三个坑

数组不可变导致的逻辑错误：我曾经在循环里用 x = x.at[i].set(...)，但忘了把结果赋值回去，结果数组根本没变。记住：.at[] 返回新数组，不是原地修改。
设备不匹配导致的性能下降：有一次我把数据放在GPU上，但计算函数没加 @jit，结果每次计算都要把数据从GPU拷回CPU，慢得离谱。加上 @jit 后，计算直接在GPU上完成，速度提升了10倍。
动态形状导致的重复编译：前面说过了，固定输入形状是 jit 高效运行的关键。如果实在需要变长输入，可以用 jax.jit(..., static_argnums=(0,)) 来指定某些参数是静态的，但这样会牺牲一些性能。

总结一下：

jnp.array 是不可变的，用 .at[] 更新。
设备内存管理要显式控制，避免数据在不同设备间频繁传输。
jit 能大幅提升性能，但要求纯函数和固定输入形状。

好了，这一章的内容就到这里。记住这三个核心概念，后面的分布式计算、自动微分、向量化映射，都是建立在这之上的。你想想看，如果连数组都搞不定，后面怎么玩分布式？

📚 课程目录

3、JAX核心：jnp.array、设备内存管理、即时编译（jit）入门

3.1 jnp.array：不只是NumPy的“复制品”

3.2 设备内存管理：数据到底在哪？

3.3 即时编译（jit）入门：让代码飞起来

3.4 知识体系图：三者如何协作？

3.5 避坑指南：我踩过的三个坑