jax.block_until_ready

jax.block_until_ready(x)

嘗試在 pytree 葉節點上調用 block_until_ready 方法。

參數:

x – 一個 pytree，通常在其葉節點上至少有一些 JAX 陣列實例。

返回:

一個 pytree，具有與輸入相同的結構和值，其中所有 JAX 陣列葉節點的值都已準備就緒。