{"record":{"id":"5b27a2840b76c8a5","repo":"jax-ml/jax","slug":"arguments-to-after-all-must-be-tokens","errorCode":null,"errorMessage":"Arguments to after_all must be tokens","messagePattern":"Arguments to after_all must be tokens","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/lax.py","lineNumber":9182,"sourceCode":"create_token_p.def_impl(partial(dispatch.apply_primitive, create_token_p))\ncreate_token_p.def_abstract_eval(lambda *_: abstract_token)\n\ndef _create_token_lowering(ctx, *operands):\n  aval_out, = ctx.avals_out\n  return [hlo.create_token()]\nmlir.register_lowering(create_token_p, _create_token_lowering)\n\n\ndef after_all(*operands):\n  \"\"\"Merges one or more XLA token values. Experimental.\n\n  Wraps the XLA after all operator.\"\"\"\n  operands = core.auto_insert_reshard(*operands)\n  return after_all_p.bind(*operands)\n\ndef _after_all_abstract_eval(*operands):\n  if any(x is not abstract_token for x in operands):\n    raise TypeError(\"Arguments to after_all must be tokens\")\n  return abstract_token\n\n\nafter_all_p = Primitive(\"after_all\")\nafter_all_p.def_impl(partial(dispatch.apply_primitive, after_all_p))\nafter_all_p.def_abstract_eval(_after_all_abstract_eval)\n\ndef _after_all_lowering(ctx, *operands):\n  aval_out, = ctx.avals_out\n  return [hlo.after_all(operands)]\nmlir.register_lowering(after_all_p, _after_all_lowering)\n\n\ndef rng_uniform(a, b, shape):\n  \"\"\"Stateful PRNG generator. Experimental and its use is discouraged.\n\n  Returns uniformly distributed random numbers in the range [a, b). If\n  b <= a, then the result is undefined, and different implementations may","sourceCodeStart":9164,"sourceCodeEnd":9200,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/lax.py#L9164-L9200","documentation":"after_all is a low-level control-flow primitive for XLA tokens (used to order side-effecting ops like infeed/outfeed and collective communication); every operand must be an abstract token. Passing an array or any non-token raises TypeError.","triggerScenarios":"lax.after_all(tok, array) or feeding the result of a non-token op into after_all; typically only in handwritten pipelines using lax.create_token and tokens explicitly.","commonSituations":"Writing custom collectives/donated-buffer pipelines; accidentally reordering arguments so a data array lands where a token is expected; mixing up after_all with other sequencing APIs.","solutions":["Only pass tokens obtained from lax.create_token, after_all, or ops that return tokens (infeed, outfeed, collectives).","Check each argument: arguments that are arrays indicate a call-site bug; move data args to the actual data op.","Prefer higher-level APIs (jax.lax.map, pmap/shard_map collectives) which manage tokens for you."],"exampleFix":"# before\ntok2 = lax.after_all(tok, some_array)\n# after\ntok2 = lax.after_all(tok, lax.create_token())  # tokens only","handlingStrategy":"type-guard","validationCode":"from jax._src.lax.lax import abstract_token\ntoks = [t for t in operands if t is not None]\n# only pass tokens: build them with lax.create_token()\ntok = lax.after_all(lax.create_token(), lax.create_token())","typeGuard":"def all_tokens(operands):\n    from jax._src.lax.lax import abstract_token\n    return all(getattr(o, 'aval', None) is abstract_token or o is abstract_token for o in operands)","tryCatchPattern":null,"preventionTips":["Only chain tokens from create_token/after_all/side-effecting ops.","Prefer high-level APIs (shard_map collectives) that handle token ordering."],"tags":["jax","lax","tokens","control-flow","xla"],"backgroundTag":"invalid-operand-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}