Device placement decides where an array lives and where its computation runs. Put participating data on an accelerator’s memory to use GPU or TPU acceleration, especially for large neural-network training or simulations. section 3.2.3
JAX provides a small set of practical operations for inspecting and moving arrays:
- Inspect available devices with jax.devices(), jax.local_devices(), jax.devices('cpu'), or jax.device_count('gpu'). section 3.2.3
- Check an array’s current device with arr.device(). section 3.2.3
- Move an array with jax.device_put(arr, target_device), or return it to host memory with jax.device_get(arr). section 3.2.3
Moving an array between CPU and GPU
arr = jnp.array([1, 42, 31337]) arr_cpu = jax.device_put(arr, jax.devices('cpu')[0]) arr_cpu.device() # CpuDevice(id=0) arr + arr_cpu # works arr_gpu = jax.device_put(arr, jax.devices('gpu')[0]) arr_gpu + arr_cpu # ValueError: incompatible devices The error identifies the operands’ different GPU and CPU placements. Move data deliberately before combining it. section 3.2.3
Listing 3.17 Getting info about devices
import jax
jax.devices() #1
>>> [gpu(id=0)] #2
jax.local_devices() #3
>>> [gpu(id=0)]
jax.devices('cpu') #4
>>> [CpuDevice(id=0)]
jax.device_count('gpu') #5
>>> 1
Use the device-query listing as a compact reference for discovering available CPU, GPU, or other devices before placing arrays. section 3.2.3
