4、JAX基础操作：jnp.array、设备内存管理、随机数生成（PRNGKey）、向量化操作

好，咱们直接进入正题。JAX 这玩意儿，表面上看跟 NumPy 长得像亲兄弟，但骨子里完全是两码事。我刚开始接触的时候，也以为就是把 np 换成 jnp 就完事了。结果呢？踩了一堆坑。今天咱们就把这四个最基础、也最核心的操作掰开揉碎讲清楚。

4.1 jnp.array：不只是换个名字

先说说 jnp.array。你可能会想：「这不就是 NumPy 的数组吗？」嗯，表面上看确实像。但 JAX 的数组有个关键特性——它是不可变的。

import jax.numpy as jnp
import numpy as np

# NumPy 可以这样
arr_np = np.array([1, 2, 3])
arr_np[0] = 99  # 没问题

# JAX 这样会报错
arr_jax = jnp.array([1, 2, 3])
# arr_jax[0] = 99  # TypeError: '' object does not support item assignment

# 正确做法：用索引更新
arr_jax_new = arr_jax.at[0].set(99)
print(arr_jax_new)  # [99, 2, 3]

我个人习惯把 JAX 数组想象成「快照」。每次操作都生成一个新快照，旧的那个还在原地不动。这在函数式编程里是常态，但如果你刚从 NumPy 转过来，确实需要适应一下。

4.2 设备内存管理：别让你的 GPU 爆掉

说到设备内存，我得先讲个真实经历。我曾经在训练一个模型时，GPU 内存莫名其妙地涨，最后直接 OOM（内存溢出）。查了半天，发现是某个中间变量没被释放。

JAX 默认是异步执行的。什么意思呢？你写一行代码，它不会立刻在 GPU 上跑，而是先构建一个计算图，等到真正需要结果时才执行。这带来了性能优势，但也让内存管理变得有点「隐形」。

4.2.1 显式设备传输

你可以用 jax.device_put() 把数据放到指定设备上：

import jax
import jax.numpy as jnp

# 查看可用设备
print(jax.devices())  # 比如 [TFRT_CPU_0, TFRT_GPU_0]

# 把数组放到 GPU 上
data = jnp.array([1, 2, 3])
data_gpu = jax.device_put(data, jax.devices('gpu')[0])

4.2.2 内存回收的坑

我曾经犯过一个错误：在循环里不断创建中间数组，以为 Python 的垃圾回收会自动处理。结果呢？GPU 内存被吃光了。

# 不好的做法：循环里创建大量中间变量
for i in range(1000):
    x = jnp.ones((1000, 1000)) * i
    # 做一些操作...

# 好的做法：复用变量，及时清理
for i in range(1000):
    x = jnp.ones((1000, 1000)) * i
    result = do_something(x)
    # 手动释放（如果需要）
    del x
    jax.clear_caches()

4.3 随机数生成：PRNGKey 的哲学

这个可能是 JAX 里最让人困惑的概念之一。为什么不能直接用 np.random.randn()？

原因很简单：JAX 是函数式的，不允许有「隐式状态」。传统的随机数生成器内部维护一个状态，每次调用都会更新。这在并行计算里会出大问题——你没法保证可重复性。

JAX 的解决方案是：显式传递随机数状态，也就是 PRNGKey。

4.3.1 基本用法

import jax.random as random

# 创建一个 key
key = random.PRNGKey(42)  # 42 是种子

# 生成随机数
key, subkey = random.split(key)
rand_vals = random.normal(subkey, shape=(3,))
print(rand_vals)

注意看，每次生成随机数前，我都用 random.split() 把 key 拆分成两个。一个用来生成随机数，另一个留着下次用。这样保证了每次的随机数都是「新鲜」的，而且可重复。

4.3.2 我踩过的坑

我曾经在写一个强化学习算法时，不小心在多个环境中复用了同一个 key。结果所有环境生成的随机数一模一样，训练直接废了。排查了半天才发现是 key 没 split 好。

# 错误示范
key = random.PRNGKey(0)
for _ in range(10):
    # 每次都使用同一个 key
    val = random.uniform(key)  # 每次都生成一样的值！
    print(val)

# 正确示范
key = random.PRNGKey(0)
for _ in range(10):
    key, subkey = random.split(key)
    val = random.uniform(subkey)  # 每次都不一样
    print(val)

4.4 向量化操作：vmap 的魔力

最后说说向量化。你可能会说：「NumPy 也有广播机制啊，JAX 有什么特别的？」

JAX 的 vmap 可以让你自动向量化任意函数。说白了，你写一个处理单个样本的函数，vmap 自动帮你改成处理批量的版本。

4.4.1 基本用法

import jax
import jax.numpy as jnp

# 定义一个处理单个向量的函数
def process_single(x):
    return x ** 2 + 2 * x + 1

# 批量处理
batch_x = jnp.array([[1, 2, 3], [4, 5, 6]])
result = jax.vmap(process_single)(batch_x)
print(result)

你看，process_single 明明只接受一维向量，但 vmap 自动把它映射到了批量数据上。这在处理神经网络时特别有用——你只需要写前向传播的逻辑，vmap 帮你搞定 batch 维度。

4.4.2 更复杂的例子

我记得有一次做图像处理，需要对每个像素应用一个复杂的变换函数。用 for 循环太慢，手动向量化又太麻烦。最后用 vmap 嵌套，一行代码搞定：

# 假设有一个复杂的像素变换函数
def pixel_transform(r, g, b):
    # 一些复杂的计算...
    return (r * 0.3 + g * 0.59 + b * 0.11)

# 对整张图片应用
image = jnp.ones((256, 256, 3))  # 假设是 RGB 图像
# 先 vmap 行，再 vmap 列
gray_image = jax.vmap(jax.vmap(pixel_transform, in_axes=-1), in_axes=(0, 0, 0))(image[..., 0], image[..., 1], image[..., 2])

知识体系总览

说了这么多，咱们用一张图来总结一下这四个基础操作的关系：

这四个基础操作，说白了就是 JAX 的「四根柱子」。你想想看，jnp.array 是数据载体，设备内存管理是资源调度，PRNGKey 是随机性的「钥匙」，vmap 是性能加速器。把它们搞明白了，后面的自动微分、JIT 编译、神经网络构建，都会顺很多。

嗯，今天就先到这儿。记住一点：JAX 的学习曲线确实比 NumPy 陡一点，但一旦你理解了它的函数式哲学，你会发现它其实很优雅。下次咱们聊聊自动微分——那个才是 JAX 真正的杀手锏。

---

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