第4章函数式纯变换

4、函数式纯变换：jax.numpy基础、随机数生成（jax.random）的独特机制、为什么需要显式传入PRNGKey

好，咱们今天聊点硬核的。JAX 这个框架，说白了就是「披着 NumPy 外衣的函数式编程引擎」。你如果只会写 np.array，那 JAX 对你来说就是个加速器。但你要是想真正用好它，就得理解它骨子里的「纯函数」基因。

我个人习惯把 JAX 看作一个「有洁癖的 NumPy」。它不允许你偷偷改数据，不允许你搞全局随机状态。一开始我觉得这很烦，后来才发现——这恰恰是它最值钱的地方。

4.1 jax.numpy：长得像 NumPy，但内核完全不同

先看个最简单的例子。你写 jnp.array([1, 2, 3])，感觉跟 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

为什么会这样？因为 JAX 遵循函数式编程的「无副作用」原则。你想想看，如果函数内部偷偷改了全局变量，那自动微分、JIT 编译这些优化就没法做了。JAX 的做法是：你想改数组？可以，返回一个新数组。

arr_jax_new = arr_jax.at[0].set(99)  # 返回新数组，原数组不变
print(arr_jax)      # [1, 2, 3]
print(arr_jax_new)  # [99, 2, 3]

嗯，这里要注意：.at[] 是 JAX 特有的索引更新语法。我刚开始用的时候老记不住，后来发现它其实很符合直觉——「在某个位置设置新值」。

小技巧： 如果你需要频繁更新数组，可以用 jnp.where 或 jnp.clip 这类函数式操作。它们返回新数组，而且能被 JIT 编译优化。

4.2 随机数生成：JAX 最「反直觉」的设计

接下来这个点，我敢说 90% 的初学者都会踩坑。JAX 的随机数生成，跟 NumPy 完全不是一个路子。

在 NumPy 里，你写 np.random.randn()，它会悄悄更新一个全局的随机状态。下次再调用，状态变了，结果也变了。这看起来很方便，但有个致命问题：不纯。

JAX 的做法是：把随机状态（PRNGKey）显式传进去，再显式传出来。

import jax.random as jrandom

# 创建一个 PRNGKey（随机种子）
key = jrandom.PRNGKey(42)

# 生成随机数，同时得到新的 key
key, subkey = jrandom.split(key)
rand_val = jrandom.normal(subkey, shape=(3,))
print(rand_val)  # [-0.1847,  0.4967,  1.5885]

# 下次再用，必须用新的 key
key, subkey = jrandom.split(key)
rand_val2 = jrandom.normal(subkey, shape=(3,))
print(rand_val2)  # 完全不同的值

你可能会问：「每次都要 split，不麻烦吗？」

我曾经也觉得麻烦，直到有一次做强化学习实验，需要完全可复现的随机轨迹。用 NumPy 的时候，我总得小心翼翼地保存和恢复随机状态。用 JAX 之后，直接把 key 作为参数传进函数，每次运行结果一模一样——爽了。

  核心要点： JAX 的随机函数都是纯函数。同样的 key 输入，永远得到同样的输出。这保证了可复现性，也让并行计算变得安全。

4.3 为什么需要显式传入 PRNGKey？

这个问题，我当年也困惑过。咱们从三个角度来理解：

可复现性：显式传入 key，意味着你可以精确控制随机性。调试的时候，固定 key 跑一遍，出 bug 了再固定 key 跑一遍，结果完全一样。这在 NumPy 里很难做到。
并行安全：JAX 支持多设备并行（GPU/TPU）。如果随机状态是全局的，多个设备同时调用就会打架。显式传入 key，每个设备用自己的 key，互不干扰。
函数式纯洁：JAX 的 JIT 编译要求函数没有副作用。如果随机函数偷偷改了全局状态，JIT 就没法优化了。显式传入 key，让随机函数变成纯函数，JIT 才能放心大胆地编译。

我举个例子你就明白了：

from jax import jit

# 错误的做法：在 JIT 函数内部创建 key
@jit
def bad_random_func():
    key = jrandom.PRNGKey(0)  # 每次调用都创建同一个 key
    return jrandom.normal(key, shape=(3,))

# 每次调用结果都一样！因为 key 没变
print(bad_random_func())  # [-0.1847,  0.4967,  1.5885]
print(bad_random_func())  # [-0.1847,  0.4967,  1.5885]

# 正确的做法：把 key 作为参数传入
@jit
def good_random_func(key):
    key, subkey = jrandom.split(key)
    return good_random_func(subkey)  # 递归？不，这里只是演示

# 实际使用时，外部管理 key
key = jrandom.PRNGKey(42)
val1, key = good_random_func(key)
val2, key = good_random_func(key)
print(val1)  # 第一次随机值
print(val2)  # 第二次随机值，不同

避坑指南： 千万不要在 JIT 函数内部创建 PRNGKey！我曾经犯过这个错，结果调试了一整天，发现每次生成的随机数都一样。正确的做法是：在函数外部管理 key，通过参数传进去。

4.4 知识体系总览

说了这么多，咱们用一张图来总结本章的核心逻辑：

这张图把本章的三个核心模块串起来了。左边是 jax.numpy 的不可变数组，中间是 jax.random 的显式 key 管理，右边是为什么需要这样做。底部三个关键词——纯、复现、并行——就是 JAX 函数式设计的精髓。

4.5 实战建议

最后，给你几个我实际项目中的经验：

养成好习惯：每次写随机函数，都把 key 作为第一个参数。这样代码一看就懂，不会出错。
善用 split：需要多个随机操作时，先 split 出多个子 key，每个操作用一个子 key。这样互不干扰，而且可复现。
别怕麻烦：显式传 key 看起来啰嗦，但调试的时候你就知道它有多香了。我有个项目跑了三天，发现结果不对，就是因为随机状态没管理好。换成 JAX 之后，再也没出过这种问题。

一句话总结： JAX 的随机数生成，就是把「隐式的全局状态」变成了「显式的参数传递」。一开始不习惯，用久了就回不去了。

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

📚 课程目录

4、函数式纯变换：jax.numpy基础、随机数生成（jax.random）的独特机制、为什么需要显式传入PRNGKey

4.1 jax.numpy：长得像 NumPy，但内核完全不同

4.2 随机数生成：JAX 最「反直觉」的设计

4.3 为什么需要显式传入 PRNGKey？

4.4 知识体系总览

4.5 实战建议