jax.Array.copy_to_host_async

jax.Array.copy_to_host_async#

abstract Array.copy_to_host_async()[源代码][源代码]#

异步地将 Array 复制到主机。

对于存在于加速器(如GPU或TPU)上的数组,JAX可能会在主机上缓存数组的值。通常,当用户请求设备上数组的值时,这会在幕后发生,但等到请求值时才启动设备到主机的复制,要求JAX在等待复制完成时阻塞调用者。

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