jax.copy_to_host_async

jax.copy_to_host_async(x)

尝试在 pytree 叶子上调用 copy_to_host_async 方法。

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

参数:

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

返回:

一个具有与输入相同结构和值的 pytree，其中所有 JAX 数组叶子的值的宿主副本都已启动。