jax.block_until_ready

jax.block_until_ready(x)

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

参数:

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

返回:

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