2、核心数据结构：DeviceArray详解，与NumPy数组的互操作，数据在CPU/GPU间的流转

好，咱们直接进入正题。JAX 里最核心的数据结构，就是 DeviceArray。说白了，它就是 JAX 世界里用来装数据的容器。你可能会问：「这不就是 NumPy 的 ndarray 吗？」嗯，有点像，但骨子里完全不同。

我个人习惯把 DeviceArray 理解为「长了腿的 NumPy 数组」。它能在 CPU 和 GPU 之间自由穿梭，而且自带「加速光环」。今天咱们就把它的底裤扒干净。

2.1 DeviceArray 是什么？

先看个最简单的例子：

import jax.numpy as jnp
import numpy as np

# 创建一个 DeviceArray
x = jnp.array([1, 2, 3, 4, 5])
print(type(x))  # <class 'jax.Array'>

注意，JAX 2.0 之后，DeviceArray 这个类名被统一成了 jax.Array。但底层逻辑完全一样。你看到 jax.Array，心里就把它当成 DeviceArray 就行。

它和 NumPy 数组最大的区别在哪？我总结三点：

• 设备绑定：每个 DeviceArray 都明确知道自己该在哪个设备上干活（CPU 还是 GPU）
• 不可变性：创建后就不能原地修改。想改？那就新建一个
• 懒加载：很多操作不会立即执行，而是等你真正需要结果时才跑

2.2 与 NumPy 数组的互操作

说实话，JAX 团队在设计时，就是奔着「让你几乎感觉不到切换」去的。所以 DeviceArray 和 NumPy 数组的互操作非常丝滑。

2.2.1 NumPy → DeviceArray

把 NumPy 数组扔给 JAX，最简单的方式：

import numpy as np

np_array = np.array([1.0, 2.0, 3.0])
jax_array = jnp.array(np_array)
print(jax_array.device())  # 输出：cuda:0 或 cpu

这里有个坑，我曾经踩过：默认情况下，JAX 会把数据放到 GPU 上（如果你有 GPU）。但有时候你希望数据留在 CPU 上做预处理。怎么办？

from jax import device_put
import jax

# 明确指定放到 CPU
cpu_array = jax.device_put(np_array, jax.devices('cpu')[0])
print(cpu_array.device())  # 输出：cpu

2.2.2 DeviceArray → NumPy

反过来也很简单：

jax_array = jnp.array([10, 20, 30])
np_array = np.asarray(jax_array)
print(type(np_array))  # <class 'numpy.ndarray'>

或者用 .copy() 方法：

np_array = jax_array.copy()  # 返回的是 NumPy 数组

2.3 数据在 CPU/GPU 间的流转

这部分是很多初学者容易懵的地方。我画了一张图，帮你理清数据流转的逻辑：

这张图展示了三种典型的数据流转路径：

1. 创建：从 NumPy 数组创建 DeviceArray，数据默认留在当前设备
2. 传输：用 device_put() 把数据从 CPU 搬到 GPU
3. 回传：用 np.asarray() 把 GPU 上的结果拉回 CPU

2.3.1 隐式传输 vs 显式传输

JAX 有个很贴心的设计：隐式传输。什么意思？你想想看，如果你在 GPU 上做计算，但输入数据在 CPU 上，JAX 会自动帮你搬过去。

# 假设当前默认设备是 GPU
cpu_data = jax.device_put(np.array([1, 2, 3]), jax.devices('cpu')[0])
gpu_data = jnp.square(cpu_data)  # 这里 JAX 自动把 cpu_data 搬到 GPU
print(gpu_data.device())  # 输出：cuda:0

听起来很方便对吧？但我要提醒你：隐式传输有性能开销。我曾经在项目中遇到过一个诡异的性能问题，排查了半天才发现是数据在 CPU 和 GPU 之间反复横跳。每次隐式传输都要走 PCIe 总线，带宽有限，频繁传输会严重拖慢速度。

2.3.2 多 GPU 场景

如果你有多个 GPU，事情会更有趣。JAX 允许你精确控制数据放到哪个设备上：

from jax import devices

gpu_0 = devices('gpu')[0]
gpu_1 = devices('gpu')[1]

data_0 = jax.device_put(np.zeros((100, 100)), gpu_0)
data_1 = jax.device_put(np.ones((100, 100)), gpu_1)

# 跨设备计算会报错
# result = data_0 + data_1  # 这行会抛出异常

嗯，这里要注意：不同设备上的 DeviceArray 不能直接做运算。你必须先把它们搬到同一个设备上。这是 JAX 的设计哲学——避免隐式的跨设备通信，让你明确知道开销在哪。

2.4 实战：一个完整的流转示例

最后，咱们看一个完整的例子，把今天讲的东西串起来：

import jax.numpy as jnp
import numpy as np
from jax import device_put, devices

# 1. 在 CPU 上准备数据
print("=== 1. 准备数据 ===")
np_data = np.random.randn(1000, 1000).astype(np.float32)
print(f"NumPy 数组类型: {type(np_data)}")

# 2. 转为 DeviceArray（默认放到当前设备）
print("\n=== 2. 转为 DeviceArray ===")
jax_data = jnp.array(np_data)
print(f"DeviceArray 所在设备: {jax_data.device()}")

# 3. 显式搬到 GPU（如果有）
print("\n=== 3. 搬到 GPU ===")
if jax.devices('gpu'):
    gpu_data = device_put(jax_data, devices('gpu')[0])
    print(f"GPU 上的 DeviceArray: {gpu_data.device()}")
else:
    print("没有 GPU，数据留在 CPU")

# 4. 在设备上做计算
print("\n=== 4. 计算 ===")
result = jnp.dot(gpu_data, gpu_data.T)
print(f"计算结果设备: {result.device()}")

# 5. 把结果拉回 CPU
print("\n=== 5. 结果回传 ===")
np_result = np.asarray(result)
print(f"回传后的类型: {type(np_result)}")
print(f"结果形状: {np_result.shape}")

这个流程，说白了就是 JAX 数据管理的标准操作。你写多了就会发现，核心就三件事：创建、传输、计算。把这三件事理清楚，DeviceArray 就没什么秘密了。

---

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