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 数组叶节点的值的主机副本已启动。