Need a tensor on a particular accelerator? jax.device_put() copies data to a specified device and returns the new array, leaving the original unchanged. Explicit placement commits the array to that device; without placement, JAX uses its default device and keeps the result uncommitted. section 3.2.3
The practical rules are:
- Pass a device such as jax.devices('cpu')[0] or jax.devices('gpu')[0] to select placement. section 3.2.3
- Treat the functionally returned array as the placed value; the input remains unchanged. section 3.2.3
- Avoid combining committed arrays from different devices: JAX raises an incompatible-devices error for the operation. section 3.2.3
- Remember that placement is asynchronous: device_put() returns immediately, so use block_until_ready() when timing completed work. section 3.2.4
Place and compare arrays
arr = jnp.array([1, 42, 31337]) arr_cpu = jax.device_put(arr, jax.devices('cpu')[0]) arr_cpu.device() # CpuDevice(id=0) arr_gpu = jax.device_put(arr, jax.devices('gpu')[0]) arr_gpu + arr_cpu # ValueError: Received incompatible devices for jitted computation. arr + arr_cpu # Array([2, 84, 62674], dtype=int32) The uncommitted arr can combine with the committed CPU array, while the committed CPU and GPU arrays cannot. section 3.2.3
Listing 3.18 Placing data to a device
arr = jnp.array([1, 42, 31337])
arr.device()
>>> gpu(id=0) #1
arr_cpu = jax.device_put(arr, jax.devices('cpu')[0]) #2
arr_cpu.device()
>>> CpuDevice(id=0) #3
arr + arr_cpu #4
>>> Array([ 2, 84, 62674], dtype=int32) #4
arr_gpu = jax.device_put(arr, jax.devices('gpu')[0]) #5
try:
arr_gpu + arr_cpu #6
except ValueError as e:
print(e)
>>> Received incompatible devices for jitted computation.
Got argument x1 of jax.numpy.add with shape int32[3] and
device ids [0] on platform GPU and argument x2 of
jax.numpy.add with shape int32[3] and device ids [0]
on platform CPU
Inspect device placement, move an array between CPU and GPU, and observe the cross-device error in one runnable listing. section 3.2.3
