jax-ml/jax · error · TypeError

Arguments to after_all must be tokens

Error message

Arguments to after_all must be tokens

What it means

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.

Source

Thrown at jax/_src/lax/lax.py:9182

create_token_p.def_impl(partial(dispatch.apply_primitive, create_token_p))
create_token_p.def_abstract_eval(lambda *_: abstract_token)

def _create_token_lowering(ctx, *operands):
  aval_out, = ctx.avals_out
  return [hlo.create_token()]
mlir.register_lowering(create_token_p, _create_token_lowering)


def after_all(*operands):
  """Merges one or more XLA token values. Experimental.

  Wraps the XLA after all operator."""
  operands = core.auto_insert_reshard(*operands)
  return after_all_p.bind(*operands)

def _after_all_abstract_eval(*operands):
  if any(x is not abstract_token for x in operands):
    raise TypeError("Arguments to after_all must be tokens")
  return abstract_token


after_all_p = Primitive("after_all")
after_all_p.def_impl(partial(dispatch.apply_primitive, after_all_p))
after_all_p.def_abstract_eval(_after_all_abstract_eval)

def _after_all_lowering(ctx, *operands):
  aval_out, = ctx.avals_out
  return [hlo.after_all(operands)]
mlir.register_lowering(after_all_p, _after_all_lowering)


def rng_uniform(a, b, shape):
  """Stateful PRNG generator. Experimental and its use is discouraged.

  Returns uniformly distributed random numbers in the range [a, b). If
  b <= a, then the result is undefined, and different implementations may

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Only pass tokens obtained from lax.create_token, after_all, or ops that return tokens (infeed, outfeed, collectives).
2. Check each argument: arguments that are arrays indicate a call-site bug; move data args to the actual data op.
3. Prefer higher-level APIs (jax.lax.map, pmap/shard_map collectives) which manage tokens for you.

Example fix

# before
tok2 = lax.after_all(tok, some_array)
# after
tok2 = lax.after_all(tok, lax.create_token())  # tokens only

from jax._src.lax.lax import abstract_token
toks = [t for t in operands if t is not None]
# only pass tokens: build them with lax.create_token()
tok = lax.after_all(lax.create_token(), lax.create_token())

def all_tokens(operands):
    from jax._src.lax.lax import abstract_token
    return all(getattr(o, 'aval', None) is abstract_token or o is abstract_token for o in operands)

When it happens

Trigger: 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.

Common situations: 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.

Related errors

• multi-platform lowering for buffer_callback
• Value of type {type(self)} is not convertible to integer ind
• TracerBoolConversionError
• {self!r} has no `sharding`.
• Cannot use `partial_eval_jaxpr_custom` with stateful jaxprs.

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