如何在JAX多GPU环境中将变量指定到目标GPU设备

针对你遇到的JAX数组默认存储在gpu:0、需要转移到gpu:1的问题，有两种常用解决方法：

方法一：转移已有的数组到目标GPU #

使用jax.device_put()函数可将现有数组转移至指定设备：

import jax

# 获取编号为1的GPU设备
target_device = jax.local_devices()[1]

# 将数组nmp转移到gpu:1，返回新的数组实例
nmp_on_gpu1 = jax.device_put(nmp, target_device)

# 验证设备归属
print(nmp_on_gpu1.device())  # 输出: gpu:1

注意：JAX数组是不可变对象，转移操作会生成新数组，原nmp仍保留在gpu:0上。

方法二：直接在目标GPU上创建数组 #

如果无需先在默认设备生成数组，可直接在gpu:1上初始化目标数组：

import jax

target_device = jax.local_devices()[1]
# 直接在gpu:1上创建全1数组
nmp_on_gpu1 = jax.device_put(jax.numpy.ones(4), target_device)

print(nmp_on_gpu1.device())  # 输出: gpu:1

内容的提问来源于stack exchange，提问作者Peyman