jax-ml/jax · error · TypeError

async done op got {aval}, want core.AbstractFuture

Error message

async done op got {aval}, want core.AbstractFuture

What it means

Async collective 'done' primitives (e.g. all_gather_done, psum_done, reduce_scatter_done) only accept `core.AbstractFuture` avals; their abstract eval raises TypeError when fed a regular array aval. The future-start op must precede the done op.

Source

Thrown at jax/_src/lax/parallel.py:3150

# Asynchronous start abstract eval.
def _async_start_abstract_eval(sync_prim, done_fun, *args, **kwargs):
  aval, effs = sync_prim.abstract_eval(*args, **kwargs)
  return core.AbstractFuture(aval, done_fun), effs

for async_prim, sync_prim, done_p in [
    (all_gather_start_p, all_gather_p, all_gather_done_p),
    (all_gather_reduced_start_p, all_gather_reduced_p, all_gather_done_p),
    (psum_invariant_start_p, psum_invariant_p, psum_done_p),
    (reduce_scatter_start_p, reduce_scatter_p, reduce_scatter_done_p),
    (all_to_all_start_p, all_to_all_p, all_to_all_done_p),
    (ppermute_start_p, ppermute_p, ppermute_done_p),
]:
  async_prim.def_effectful_abstract_eval(
      partial(_async_start_abstract_eval, sync_prim, done_p.bind))

def _async_done_abstract_eval(aval):
  if not isinstance(aval, core.AbstractFuture):
    raise TypeError(f"async done op got {aval}, want core.AbstractFuture")
  return aval.inner_aval

for p, target in [
    (all_gather_done_p, "all-gather-done"),
    (psum_done_p, "all-reduce-done"),
    (reduce_scatter_done_p, "reduce-scatter-done"),
    (all_to_all_done_p, "all-to-all-done"),
    (ppermute_done_p, "collective-permute-done"),
]:
  p.def_abstract_eval(_async_done_abstract_eval)
  mlir.register_lowering(p, partial(_async_done_lowering, target))


mlir.register_lowering(
    reduce_scatter_start_p,
    partial(_reduce_scatter_lowering, lax.add_p, is_async=True))

mlir.register_lowering(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use the matching async start op first and feed its future result to the done op
2. Use the public async wrapper APIs rather than binding done primitives directly
3. Verify the value between start and done passes through unchanged (no intermediate ops consuming the future)

Example fix

# before
done = reduce_scatter_done_p.bind(x)  # x is a plain array
# after
fut = reduce_scatter_start(x, ...)
out = reduce_scatter_done(fut)

# only call done ops on the output of the matching start op
fut = async_start(x)
out = async_done(fut)

def is_future(x) -> bool:
    return isinstance(jax.typeof(x), jax._src.core.AbstractFuture)

When it happens

Trigger: Calling the done variant of an async collective directly on a plain array instead of on the future token returned by the corresponding async start op; mixing sync and async variants of a collective.

Common situations: Hand-composing async collectives (start/done pairs) for overlap; refactoring code that used the sync primitive and forgetting to switch the operand to the start-op output.

Related errors

• {name} was requested to map a value of non-array type {core.
• primal and tangent arguments to jax.jvp must be tuples or li
• {prim_name} takes a scalar pred as argument, got {pred}
• `compute_on`'s compute_type argument must be a string.
• Cannot interpret value of type {typ} as an abstract array; i

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/2f26ae668524e9e5. Report an issue: GitHub.