jax.Array.copy_to_host_async

Array.copy_to_host_async()

将 Array 异步复制到主机。

对于驻留在加速器（例如 GPU 或 TPU）上的数组，JAX 可能会将数组的值缓存到主机上。通常，当用户请求设备上数组的值时，此过程会在后台进行，但等到请求值时才启动设备到主机复制，这需要 JAX 阻塞调用者，直到复制完成。

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