JAX基础入门：张量操作、设备管理与随机数生成

大家好，欢迎来到第二讲。今天咱们聊聊JAX最核心的三个基本功：张量怎么玩、设备怎么管、随机数怎么生成。这三个东西，说白了就是JAX的「三驾马车」。我在做RL项目时，80%的bug都出在这三个地方。尤其是随机数管理，嗯，这个坑我踩过不止一次。

一、张量操作：从NumPy到JAX的无缝迁移

如果你用过NumPy，那JAX的张量操作你基本已经会了。JAX的API设计哲学就是「像NumPy一样，但更快、更自动」。我个人习惯把JAX的数组叫做「张量」，虽然它跟PyTorch的张量不完全一样。

来看个例子：

import jax
import jax.numpy as jnp

# 创建张量
x = jnp.array([1, 2, 3, 4, 5])
y = jnp.ones((3, 4))
z = jnp.zeros((2, 3, 4))

# 基本运算
a = jnp.sin(x) + jnp.cos(y)
b = jnp.dot(x, x)  # 内积
c = jnp.matmul(y, y.T)  # 矩阵乘法

# 索引与切片（注意！不能直接赋值）
print(x[1:3])  # 没问题
# x[0] = 10   # 会报错！

为什么会这样？因为JAX追求纯函数式编程。每次操作都返回新数组，不修改原数据。这在自动微分和JIT编译时特别有用。我在做策略梯度算法时，就因为忘了这个特性，debug了整整一下午。

二、设备管理：CPU/GPU/TPU的「交通指挥」

JAX最让我惊艳的一点，就是设备管理。你写一份代码，它自动跑在CPU、GPU或TPU上。不需要像PyTorch那样手动 .to(device)。说白了，JAX帮你做了「设备无关」的抽象。

但这里有个坑：JAX默认不会自动把数据放到GPU上。你得显式指定设备，或者用 jax.devices() 查看可用设备。

import jax

# 查看当前可用设备
devices = jax.devices()
print(f"可用设备: {devices}")

# 手动指定设备
with jax.default_device(jax.devices('gpu')[0]):
    x = jnp.array([1, 2, 3])
    print(x.device)  # 输出: gpu:0

# 或者用 jax.device_put 显式放置
y = jax.device_put(jnp.array([4, 5, 6]), jax.devices('cpu')[0])

我在项目中遇到过一个问题：在TPU上跑RL训练，结果某个张量莫名其妙在CPU上计算，导致训练慢得像蜗牛。后来发现是某个自定义函数没有用JIT编译。所以我的建议是：能用 @jit 的地方尽量用，它会自动帮你做设备管理。

设备类型	适用场景	注意事项
CPU	小规模调试、数据预处理	速度慢，但兼容性好
GPU	RL训练、神经网络推理	显存有限，注意OOM
TPU	大规模并行计算、矩阵运算	需要特定配置，不适合小batch

三、随机数生成：Key管理机制的精髓

这是JAX最独特、也最容易让人困惑的地方。传统的随机数生成器（比如NumPy）是「有状态的」——你调用一次 np.random.randn()，内部状态就变了。但JAX是纯函数式的，它不允许有隐藏状态。所以JAX搞了个「显式Key管理」机制。

说白了，你得自己管理一个「随机数钥匙」。每次生成随机数，都要用这个Key，然后生成一个新的Key。

import jax.random as random

# 初始化一个Key
key = random.PRNGKey(42)

# 用Key生成随机数，同时得到新Key
key, subkey = random.split(key)
x = random.normal(subkey, shape=(3, 4))

# 再生成一批
key, subkey = random.split(key)
y = random.uniform(subkey, shape=(2, 3), minval=0.0, maxval=1.0)

# 如果你需要多个独立随机流
keys = random.split(key, 5)  # 生成5个独立Key

你想想看，为什么JAX要这么设计？因为可复现性。在RL中，你经常需要复现某个随机种子下的实验结果。用JAX的Key机制，你可以精确控制每一次随机操作。我在做PPO算法时，就靠这个机制复现了一个bug——原来是我的动作采样和策略更新用了同一个Key，导致梯度计算错误。

知识体系总览

下面这张图，是我自己总结的JAX基础三件套的关系。你看一眼就能明白它们怎么配合：

你看，张量操作是基础数据载体，设备管理决定计算发生在哪里，随机数生成则提供不确定性。三者配合，才能构建出完整的RL训练流程。我个人觉得，理解这个关系比死记API更重要。

实战小练习

最后，给你留个小练习。试试用JAX实现一个简单的「经验回放缓冲区」的采样逻辑：

def sample_experience(key, buffer, batch_size):
    """
    从经验回放缓冲区中随机采样batch_size条经验
    buffer: 形状为 (buffer_size, experience_dim) 的数组
    """
    # 你的代码应该：
    # 1. 用key生成随机索引
    # 2. 从buffer中取出对应行
    # 3. 返回采样结果和新key
    pass

提示：用 random.randint 生成索引，注意不要重复使用同一个key。

好了，这一讲就到这里。记住：张量操作要函数式，设备管理要显式，随机数Key要分叉。这三个习惯养成了，后面学JIT编译和自动微分会轻松很多。

---