jax.block_until_ready

jax.block_until_ready#

jax.block_until_ready(x)[source]#

Tries to call a block_until_ready method on pytree leaves.

Parameters:

x – a pytree, usually with at least some JAX array instances at its leaves.

Returns:

A pytree with the same structure and values of the input, where the values of all JAX array leaves are ready.