jax.block_until_ready

jax.block_until_ready(x)

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

参数:

x – 一个 pytree，通常其叶子中至少有一些 JAX 数组实例。

返回:

一个具有与输入相同的结构和值的 pytree，其中所有 JAX 数组叶子的值都已准备就绪。