2、JAX基础入门：jnp.array与np.array的区别、设备内存管理、随机数生成、函数式纯化原则

好，咱们正式开始JAX的实战之旅。这一节我打算把JAX最基础、也最容易踩坑的几个点讲透。你可能会觉得「不就是数组操作吗，能有什么花样？」——嗯，我当初也是这么想的，结果第一次写JAX代码就翻车了。

2.1 jnp.array vs np.array：不只是换个名字

先看个最直观的例子。你用NumPy写习惯了，可能会觉得JAX的数组就是换个包名。其实不然。

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

# 看起来一样？
a_np = np.array([1, 2, 3])
a_jnp = jnp.array([1, 2, 3])

print(type(a_np))   # <class 'numpy.ndarray'>
print(type(a_jnp))  # <class 'jax.Array'>

类型不同只是表象。真正的区别在于：jnp.array是不可变的。你没看错，JAX的数组一旦创建就不能原地修改。

a = jnp.array([1, 2, 3])
a[0] = 100  # TypeError: JAX arrays are immutable

我刚开始用JAX时，习惯性地写了个 arr[idx] = new_val，结果程序直接崩了。后来才意识到，JAX的设计哲学就是「函数式」，所有操作都返回新数组，不修改原数据。

另一个关键区别是数据类型默认行为。NumPy默认用64位，JAX默认用32位。这在精度敏感的场景下要特别注意。

print(np.array([1.0]).dtype)   # float64
print(jnp.array([1.0]).dtype)  # float32

我个人习惯在写JAX代码时，显式指定dtype，避免精度问题带来的bug。

2.2 设备内存管理：.device_buffer 的秘密

JAX的数组默认放在哪里？答案是：GPU或TPU上。如果你有加速器，JAX会自动把数据放上去。这和NumPy完全不一样——NumPy的数据永远在CPU内存里。

看个例子：

import jax
import jax.numpy as jnp

# 检查默认设备
print(jax.devices())  # 比如 [TFRT_CPU_0] 或 [gpu:0]

# 创建数组
x = jnp.array([1, 2, 3])
print(x.device())  # 显示在哪个设备上

# 查看底层缓冲区
print(x.device_buffer)  # <DeviceArray at 0x...>

.device_buffer 这个属性，说白了就是JAX数组在设备上的「真实地址」。你想想看，如果数据在GPU上，你直接拿CPU代码去读，肯定出问题。

我曾经在调试一个分布式训练任务时，发现某个中间结果死活不对。查了半天，原来是两个数组一个在GPU:0上，一个在GPU:1上，操作时触发了隐式数据传输。嗯，从那以后我养成了习惯：每次创建数组都检查一下 .device()。

2.3 随机数生成：jax.random 的正确打开方式

JAX的随机数生成和NumPy完全不同。NumPy里你写 np.random.randn() 就行，JAX不行——它需要你显式传入一个随机数状态（key）。

import jax.random as jrandom

# 错误写法（别学我）
# x = jrandom.normal(shape=(3,))  # TypeError: missing required argument: key

# 正确写法
key = jrandom.PRNGKey(42)
x = jrandom.normal(key, shape=(3,))
print(x)

为什么要这么麻烦？因为JAX要求纯函数——同样的输入必须产生同样的输出。如果随机数生成器内部维护一个隐式状态，那就破坏了纯函数原则。

这里有个坑：不要重复使用同一个key。

key = jrandom.PRNGKey(42)

# 错误：两次用同一个key，结果一样
a = jrandom.normal(key, shape=(3,))
b = jrandom.normal(key, shape=(3,))
print(a == b)  # True

# 正确：用 split 分裂出新的key
key, subkey = jrandom.split(key)
a = jrandom.normal(key, shape=(3,))
b = jrandom.normal(subkey, shape=(3,))
print(a == b)  # False

我个人习惯在函数开头先 key, *subkeys = jrandom.split(key, num=N)，一次性分好所有子key，后面直接用。这样既清晰又安全。

2.4 函数式纯化原则：JAX的灵魂

前面其实一直在提「纯函数」。到底什么是JAX的纯函数原则？说白了就两条：

1. 无副作用：函数不修改外部状态
2. 引用透明：相同输入永远得到相同输出

看个对比：

# NumPy风格（有副作用）
def add_one_np(arr):
    arr += 1  # 修改了原数组
    return arr

# JAX风格（纯函数）
def add_one_jax(arr):
    return arr + 1  # 返回新数组，原数组不变

你可能会问：「这有什么好处？」好处太大了。JAX的 jit、vmap、grad 等所有高级功能，都依赖这个纯函数假设。如果你写了不纯的函数，JAX编译出来的结果可能完全不对。

我记得有次帮同事调试一个梯度计算错误的问题。他的损失函数里用了一个全局计数器来记录训练步数——这在PyTorch里完全没问题，但在JAX里，grad 会把这个计数器当作常量处理，导致梯度完全算错。改掉这个「不纯」的习惯后，问题立刻解决了。

2.5 知识体系总览

下面这张图把这一节的核心逻辑串起来了。你可以看到，所有知识点都围绕「函数式纯化」这个核心展开。

从这张图可以看得很清楚：jnp.array的不可变性、设备内存的显式管理、随机数生成器的状态传递，这三者都是「函数式纯化」这个核心原则的具体体现。你只要抓住这个核心，JAX的其他高级特性（jit、vmap、grad等）理解起来就会顺畅很多。

好了，这一节的内容就到这里。记住这些基础，后面我们写自定义算子时，你会感谢自己现在打下的基础。