Committed data is a JAX array explicitly placed on a device, making it sticky to that device. Unlike uncommitted data, it does not move or reshard automatically. section 3.2.3
Use jax.device_put() with a device argument to commit an array. Check placement with device(); mixing committed arrays on incompatible devices can raise a device error. 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
Compare GPU- and CPU-committed arrays, then observe successful addition on compatible placement and an error when devices conflict. section 3.2.3
