jax.Array.copy_to_host_async

Array.copy_to_host_async()

将 Array 异步复制到主机。

对于存在于加速器（如 GPU 或 TPU）上的数组，JAX 可能会在主机上缓存数组的值。通常情况下，当用户请求设备上数组的值时，这会在后台发生。但是，等到值被请求时才发起设备到主机的复制操作，需要 JAX 在等待复制完成时阻塞调用者。

copy_to_host_async 请求 JAX 填充其主机上的数组缓存，但不会等待复制完成。这可能会加速将来对数组内容的本地访问。