2. JAX核心数据结构：DeviceArray、PRNGKey随机数生成、数组创建与操作
好，咱们直接进入正题。JAX 的核心，说白了就是围绕几个关键数据结构转的。你如果搞懂了 DeviceArray 和 PRNGKey，那 JAX 的感知管线你就拿下一半了。我刚开始从 NumPy 切到 JAX 时，也踩过不少坑，今天咱们一个一个捋清楚。
2.1 DeviceArray：JAX 的“亲儿子”数组
JAX 里的数组，不叫 ndarray，叫 DeviceArray。名字已经告诉你答案了——它活在设备上（GPU/TPU）。
你想想看，在 NumPy 里你创建一个数组，数据就在 CPU 内存里。但在 JAX 里，默认情况下，数组直接扔到 GPU 上。这意味着什么？意味着你没法直接 print 看它的值，得先把它拉回 CPU。
核心区别一句话：NumPy 数组在 CPU 上，JAX 的 DeviceArray 在加速器上。
import jax
import jax.numpy as jnp
# 创建一个 DeviceArray
arr = jnp.array([1.0, 2.0, 3.0])
print(type(arr)) # <class 'jaxlib.xla_extension.ArrayImpl'>
# 想看看值？用 .block_until_ready() 或者直接 print
print(arr) # 会自动同步，但最好养成显式同步的习惯我个人习惯在调试时用 jax.device_get() 把数组拉回 NumPy，这样能避免很多隐式同步带来的性能问题。
小技巧：如果你想把 DeviceArray 转成 NumPy 数组，用 np.array(jax_array) 或者 jax_array.__array__()。但注意，这会触发一次设备到主机的拷贝，别在热循环里这么干。
2.2 PRNGKey：JAX 的随机数生成，跟你想的不一样
JAX 的随机数生成，可能是初学者最头疼的地方。为什么？因为它没有全局随机状态。
在 NumPy 里，你调 np.random.randn()，内部状态自动更新。但在 JAX 里，所有随机函数都需要你显式传入一个 PRNGKey。这设计一开始我也觉得烦，但后来发现——这其实是好事。
为什么？因为 JAX 的函数式编程范式要求“纯函数”——同样的输入，永远给出同样的输出。如果随机数生成器偷偷改了内部状态，那函数就不纯了。
from jax import random
# 创建一个 PRNGKey
key = random.PRNGKey(42) # 42 是种子，随便选
# 生成随机数
key, subkey = random.split(key) # 分裂 key，得到新的 key 和子 key
rand_vals = random.normal(subkey, shape=(3,))
print(rand_vals)我曾经在项目中犯过一个低级错误：在循环里反复用同一个 key 生成随机数，结果每次迭代得到的随机序列完全一样。嗯，那模型训练直接废了。
避坑指南：千万不要重复使用同一个 key！每次生成随机数前，先用 random.split() 分裂出新的 key。我习惯在训练循环开头统一 split 一批 key，然后按需取用。
你可能会问：为什么不用 random.fold()？fold 是把 key 和一个整数结合生成新 key，适合在循环里按索引生成。我个人更推荐 split，因为它更直观，不容易出错。
2.3 数组创建与操作：从 NumPy 到 JAX 的无缝迁移
好消息是，JAX 的数组操作 API 几乎和 NumPy 一模一样。你如果会 NumPy，那 JAX 上手基本零成本。
但有几个关键区别，我建议你记牢：
| 操作 | NumPy | JAX | 注意点 |
|---|---|---|---|
| 创建全零数组 | np.zeros((3,4)) |
jnp.zeros((3,4)) |
完全一样 |
| 创建全一数组 | np.ones((2,3)) |
jnp.ones((2,3)) |
完全一样 |
| 创建单位矩阵 | np.eye(5) |
jnp.eye(5) |
完全一样 |
| 创建等差数列 | np.arange(0, 10, 2) |
jnp.arange(0, 10, 2) |
完全一样 |
| 创建均匀分布 | np.linspace(0, 1, 5) |
jnp.linspace(0, 1, 5) |
完全一样 |
| 重塑形状 | arr.reshape(2, -1) |
arr.reshape(2, -1) |
完全一样 |
| 转置 | arr.T |
arr.T |
完全一样 |
| 拼接 | np.concatenate([a,b]) |
jnp.concatenate([a,b]) |
完全一样 |
你看，基本就是换个前缀的事。但有一个地方要特别小心——原地修改。
重要提醒：JAX 的数组是不可变的！你不能写 arr[0] = 1，会报错。如果你需要“修改”数组，实际上是在创建新数组。这在自动驾驶感知中处理点云数据时尤其要注意——别想着原地更新，用 jnp.where 或 jnp.scatter_add 这类函数式操作。
# 错误示范：原地修改
# arr[0] = 100 # 会报错！
# 正确做法：用索引更新
arr = jnp.array([1, 2, 3])
new_arr = arr.at[0].set(100) # 返回新数组，原数组不变
print(new_arr) # [100, 2, 3]这个 .at[] 语法，我一开始也觉得别扭。但用习惯了你会发现，它其实更安全——你永远不会不小心改了原始数据。在自动驾驶感知中，我们经常要同时维护多个数据副本（比如原始点云、降采样后的点云、体素化后的点云），不可变性让这一切变得清晰可控。
2.4 知识体系总览
下面这张图，帮你把本章的核心逻辑串起来：
这张图把三个核心模块的关系画得很清楚。你从左边 DeviceArray 出发，经过 PRNGKey 的随机数生成，再到数组的各种创建和操作，最后都落到同一个原则：函数式、不可变、显式状态。记住这三点，JAX 的坑你就能避开一大半。
我的建议：刚开始用 JAX 时，别急着上复杂模型。先花一天时间，把 jnp 的各种操作和 random 模块玩熟。我在做激光雷达点云处理时，就是靠这些基础操作一步步搭起来的。基础不牢，后面写 JIT 编译和 vmap 时会很痛苦。
公众号:蓝海资料掘金营，微信deep3321