jax.copy_to_host_async

jax.copy_to_host_async(x)

尝试在 Py树 的叶子上调用 copy_to_host_async 方法。

对于每个叶子，此方法将尝试调用其上的 copy_to_host_async 方法。如果叶子不是 JAX 数组，或者叶子没有 copy_to_host_async 方法，则此方法不会对该叶子执行任何操作。

参数:

x – 一个 Py树，通常其叶子至少包含一些 JAX 数组实例。

返回:

一个 Py树，其结构和值与输入相同，其中所有 JAX 数组叶子的值的主机拷贝已开始。