jax.block_until_ready

jax.block_until_ready(x)

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

参数：

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

返回值：

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