jax.block_until_ready

jax.block_until_ready(x)

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

参数:

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

返回:

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