Getting Started
Reference Documentation
Advanced JAX Tutorials
Notes
custom_vjp
nondiff_argnums
Developer documentation
API documentation
Tries to call a block_until_ready method on pytree leaves.
block_until_ready
x – a pytree, usually with at least some JAX array instances at its leaves.
A pytree with the same structure and values of the input, where the values of all JAX array leaves are ready.