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.