{"record":{"id":"2f26ae668524e9e5","repo":"jax-ml/jax","slug":"async-done-op-got-aval-want-core-abstractfuture","errorCode":null,"errorMessage":"async done op got {aval}, want core.AbstractFuture","messagePattern":"async done op got (.+?), want core\\.AbstractFuture","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/parallel.py","lineNumber":3150,"sourceCode":"# Asynchronous start abstract eval.\ndef _async_start_abstract_eval(sync_prim, done_fun, *args, **kwargs):\n  aval, effs = sync_prim.abstract_eval(*args, **kwargs)\n  return core.AbstractFuture(aval, done_fun), effs\n\nfor async_prim, sync_prim, done_p in [\n    (all_gather_start_p, all_gather_p, all_gather_done_p),\n    (all_gather_reduced_start_p, all_gather_reduced_p, all_gather_done_p),\n    (psum_invariant_start_p, psum_invariant_p, psum_done_p),\n    (reduce_scatter_start_p, reduce_scatter_p, reduce_scatter_done_p),\n    (all_to_all_start_p, all_to_all_p, all_to_all_done_p),\n    (ppermute_start_p, ppermute_p, ppermute_done_p),\n]:\n  async_prim.def_effectful_abstract_eval(\n      partial(_async_start_abstract_eval, sync_prim, done_p.bind))\n\ndef _async_done_abstract_eval(aval):\n  if not isinstance(aval, core.AbstractFuture):\n    raise TypeError(f\"async done op got {aval}, want core.AbstractFuture\")\n  return aval.inner_aval\n\nfor p, target in [\n    (all_gather_done_p, \"all-gather-done\"),\n    (psum_done_p, \"all-reduce-done\"),\n    (reduce_scatter_done_p, \"reduce-scatter-done\"),\n    (all_to_all_done_p, \"all-to-all-done\"),\n    (ppermute_done_p, \"collective-permute-done\"),\n]:\n  p.def_abstract_eval(_async_done_abstract_eval)\n  mlir.register_lowering(p, partial(_async_done_lowering, target))\n\n\nmlir.register_lowering(\n    reduce_scatter_start_p,\n    partial(_reduce_scatter_lowering, lax.add_p, is_async=True))\n\nmlir.register_lowering(","sourceCodeStart":3132,"sourceCodeEnd":3168,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/parallel.py#L3132-L3168","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Use the matching async start op first and feed its future result to the done op","Use the public async wrapper APIs rather than binding done primitives directly","Verify the value between start and done passes through unchanged (no intermediate ops consuming the future)"],"exampleFix":"# before\ndone = reduce_scatter_done_p.bind(x)  # x is a plain array\n# after\nfut = reduce_scatter_start(x, ...)\nout = reduce_scatter_done(fut)","handlingStrategy":"type-guard","validationCode":"# only call done ops on the output of the matching start op\nfut = async_start(x)\nout = async_done(fut)","typeGuard":"def is_future(x) -> bool:\n    return isinstance(jax.typeof(x), jax._src.core.AbstractFuture)","tryCatchPattern":null,"preventionTips":["Never feed plain arrays to *_done primitives; always pair start->done"],"tags":["jax","async-collectives","future","type-error"],"backgroundTag":"wrong-operand-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}