jax.block_until_ready

jax.block_until_ready(x)

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

参数：

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

返回：

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