jax.Array.copy_to_host_async#

Array.copy_to_host_async()[source]#

Copies anArray to the host asynchronously.

For arrays that live an an accelerator, such as a GPU or a TPU, JAX maycache the value of the array on the host. Normally this happensbehind the scenes when the value of an on-device array is requested by theuser, but waiting to initiate a device-to-host copy until the value isrequested requires that JAX block the caller while waiting for the copy tocomplete.

copy_to_host_async requests that JAX populate its on-host cache of anarray, but does not wait for the copy to complete. This may speed up afuture on-host access to the array’s contents.