block_until_ready for Sync

Wait for the result you are measuring

A JAX call can return an array before the device has finished computing its contents. Python can then submit more work while earlier work runs. Reading an array's shape does not require its numerical values to be ready.

This matters when timing a calculation: stopping the clock as soon as the call returns can miss unfinished work. It does not always measure only dispatch, since tracing, compilation or other work may also occur. There is no fixed ratio between dispatch time and execution time.

result.block_until_ready() waits for that array's value to be ready and returns the array. It does not convert the result into a Python number or copy its contents to a NumPy array.

Time a warmed-up calculation

import time
import jax
import jax.numpy as jnp

@jax.jit
def sum_sq(x):
    return jnp.sum(x ** 2)

x = jnp.ones((10_000,), dtype=jnp.float32)
x.block_until_ready()            # finish input setup
sum_sq(x).block_until_ready()    # compile and finish a warm-up call

start = time.perf_counter()
result = sum_sq(x)
result.block_until_ready()       # wait before stopping the clock
elapsed = time.perf_counter() - start
print(result, elapsed)           # result is 10000; time varies

This measures elapsed host time for dispatch, execution and waiting after warm-up. It is not a measurement of device kernel time alone. Repeat measurements to assess variability. To study startup instead, measure the first call separately and say that it includes compilation.

Wait on the final output whose completion you need. Waiting on an intermediate array does not establish that later calculations finished. For several array outputs, jax.block_until_ready(outputs) can wait on the array leaves of a tuple or other pytree. Neither form is a blanket guarantee that unrelated device work or callbacks have finished.

Keep the wait outside JIT

Synchronization is a host action. The array method is not available on a tracer, so placing result.block_until_ready() inside a jitted body can raise AttributeError during tracing. Compile the numerical helper, call it from Python, then wait on its returned array as shown above.

JAX already tracks dependencies between its array operations. You do not need to block between every step to get the correct result. Reading actual values in Python, such as with float(result), also waits as needed, but changes the operation being measured. Explicit waits are useful when you need a completion boundary, including timing and profiling, not only when you need a Python value.