What blocking actually waits on

Open the method and there is less machinery than the name suggests. It checks the array has not been deleted, walks its list of device buffers, blocks on each one, and returns the array itself. No device queue is drained, no other computation is waited on, and nothing is copied to the host.

Returning self is what makes the harness idiom work. jax.block_until_ready(f(*args)) reads as one expression because the call hands the array back, so a warmup line can block and bind on one line. Chapter 11 owns the harness; the reason the idiom composes is this line.

The buffer list is the part that generalizes. On one device it has one entry, so the loop runs once. On a sharded array it has one entry per addressable device, and the call returns when the slowest of them is ready, which is the sense in which it waits for a computation and not for a device.

  @use_cpp_method()
  def block_until_ready(self):
    self._check_if_deleted()
    for db in self._arrays:
      db.block_until_ready()
    return self

Every array carries a non-blocking poll beside the blocking one. is_ready() returns whether the buffers are done, right now, and never waits, which makes asynchronous dispatch something you can watch rather than something you infer from a stopwatch.

Dispatch a large matmul on an already-compiled function and the poll says False on the line right after it. Block, and it says True. Two booleans, no timer, and the observation does not vary with what else the machine happens to be doing, which a timing number would.

Compilation has to be out of the way first or the poll measures the wrong thing. The warmup call in the snippet is there for that reason, and chapter 11 is where compilation-in-the-measurement gets its full treatment.

import jax
import jax.numpy as jnp

mm = jax.jit(lambda a: a @ a)
x = jnp.ones((3000, 3000))
mm(x).block_until_ready()      # compile, then start clean

y = mm(x)
print(y.is_ready())
print(y.block_until_ready() is y)
print(y.is_ready())

# False
# True
# True

Dispatch a small computation and then a large one, and block on the small one's output. The large one is still running when the call returns, and its poll says so. Nothing about blocking on an array drains the work queued behind it.

That is worth checking against the reverse order, because the reverse order looks like the same experiment and is not. Dispatch the large one first, block on the small one's output, and the large one has finished too, since a single CPU device runs its queue in dispatch order. The independence you can observe here is one-directional, and reading only the second case would teach the wrong rule.

For a benchmark the consequence is direct. If a step produces several outputs and the harness blocks on one of them, it has timed the path to that one output, not the step. Blocking the whole pytree of results is the version that measures what the loop actually does.

import jax
import jax.numpy as jnp

mm = jax.jit(lambda a: a @ a)
big, small = jnp.ones((3000, 3000)), jnp.ones((64, 64))
mm(big).block_until_ready()
mm(small).block_until_ready()

s = mm(small)                  # dispatched first
b = mm(big)                    # dispatched second
print(s.block_until_ready() is s)
print(b.is_ready())
print(b.block_until_ready().is_ready())

# True
# False
# True

The module-level jax.block_until_ready is the version to use on a step's outputs, and it does three things the method cannot. It flattens a pytree, it batches the arrays into a single runtime call rather than blocking them one at a time, and it hands the original object back so the call sits anywhere in an expression.

The skipping is where care is needed. Any leaf without a block_until_ready method is passed to a helper that catches AttributeError and returns it unchanged, so NumPy arrays, Python floats, and strings pass through unchanged and unmentioned. Hand it a pytree with no JAX arrays in it at all and the function does nothing and reports nothing, which in a harness reads exactly like a successful block.

The batched path is a real difference for a step with many outputs. One batched_block_until_ready call across the whole list beats a Python loop calling the method per array, and defaulting to the module-level function is how you get that without deciding it at each call site.

def block_until_ready(x):
  """Tries to call a ``block_until_ready`` method on pytree leaves."""
  def try_to_block(x):
    try:
      return x.block_until_ready()
    except AttributeError:
      return x

  arrays = []
  for leaf in tree_leaves(x):
    if isinstance(leaf, array.ArrayImpl):
      arrays.append(leaf)
    else:
      try_to_block(leaf)

  if not arrays:
    # `arrays` will be empty if tree_leaves(x) is empty or all leaves are not
    # jax.Array.
    pass
  elif len(arrays) == 1:
    # Fast path for single array.
    try_to_block(arrays[0])
  else:
    # Optimized for multiple arrays.
    xc.batched_block_until_ready(arrays)

  return x

Put a block inside a jitted function and it never runs as a block. Tracer defines block_until_ready as a property that raises AttributeError, and the comment above it says why: raising that particular exception keeps hasattr and getattr checks working the way callers expect.

Which means the module-level wrapper, whose helper catches exactly AttributeError, treats a tracer as an unblockable leaf and moves on without a word. A jax.block_until_ready call inside a jitted function is a no-op with no diagnostic, and it is a plausible thing to write when a harness gets refactored and a block ends up on the wrong side of the decorator.

The second no-op is the one from the previous section, a pytree that turned out to hold no JAX arrays. Both fail the same way, quietly and while looking correct, so a harness worth trusting asserts on the arrays it meant to block rather than assuming the call found them.

A block that found nothing to block looks exactly like a block that worked.

  @property
  def block_until_ready(self):
    # Raise AttributeError for backward compatibility with hasattr() and getattr() checks.
    raise AttributeError(self,
      f"The 'block_until_ready' method is not available on {self._error_repr()}."
      f"{self._origin_msg()}")

# >>> params = {"w": jnp.ones((2, 2)), "cached": np.ones(4), "lr": 0.01}
# >>> out = jax.block_until_ready(params)
# >>> out is params, out["w"] is params["w"]
# (True, True)
#
# >>> @jax.jit
# ... def step(v):
# ...     try:
# ...         v.block_until_ready()
# ...     except AttributeError:
# ...         print("a tracer has no block_until_ready")
# ...     return v * 2
# >>> step(jnp.ones(3)).block_until_ready()
# a tracer has no block_until_ready

Side effects have their own wait. jax.effects_barrier() is one line, and it blocks on the runtime's effect tokens rather than on any array, so it is the call that waits for a debug.print or an io_callback to have actually happened. Blocking on a function's output array says nothing about them, because the output buffer and the effect token are separate things to be ready.

Nor does blocking fetch anything. block_until_ready leaves the values where they are; np.asarray(x) is the call that moves them to the host, and copy_to_host_async is the call that starts that move without waiting for it. On this CPU backend the difference costs nothing, since the buffer is already host memory, so the separation is easiest to keep straight by reading the three call sites rather than by timing them here.

Three waits, then, with three different meanings. One array's buffers, the whole runtime's pending effects, and a transfer to the host. Using the first when you needed one of the other two is what produces a benchmark or a log that is wrong without saying so.

def effects_barrier():
  """Waits until existing functions have completed any side-effects."""
  dispatch.runtime_tokens.block_until_ready()