jax-ml/jax · error · TypeError

scan body output must be a pair, got {}.

Error message

scan body output must be a pair, got {}.

What it means

The scan body function must return exactly a pair (carry, ys). If it returns a single value, a 3-tuple, or an unpackable pytree of wrong length, this TypeError is raised after tracing.

Source

Thrown at jax/_src/lax/control_flow/loops.py:384

    for i in maybe_reversed(range(length)):
      xs_slice = xs_flat.map(lambda x: slicing.index_in_dim(x, i, keepdims=False))
      carry, y = f(carry, xs_slice.unflatten())
      ys.append(y)
    stack = lambda *ys: _stack(ys)
    stacked_y = tree_map(stack, *maybe_reversed(ys))
    return carry, stacked_y

  if config.mutable_array_checks.value:
    check_no_aliased_ref_args(lambda: dbg_body, list(args_avals), list(args))

  x_avals = xs_avals.map(lambda aval: core.mapped_leading_aval(length, aval))
  def _create_jaxpr(carry_avals):
    new_arg_avals = ft.pack(((carry_avals, x_avals), {}))
    jaxpr, out_avals = pe.trace_to_jaxpr(f, new_arg_avals, dbg_body)
    jaxpr, consts = pe.separate_consts(jaxpr)
    if not out_avals.unpackable or len(out_avals.unpack()) != 2:
      msg = "scan body output must be a pair, got {}."
      raise TypeError(msg.format(out_avals.unflatten()))
    return jaxpr, out_avals, consts

  # The carry input and output avals must match exactly. However, we want to account for
  # the case when init contains weakly-typed values (e.g. Python scalars), with avals that
  # may not match the output despite being compatible by virtue of their weak type.
  # To do this, we compute the jaxpr in two passes: first with the raw inputs, and if
  # necessary, a second time with modified init values.
  # TODO(dougalm): this two-pass stuff is expensive (exponential in scan nesting
  # depth) and incomplete (because in the general case it takes more than two passes).
  # Let's get rid of it, perhaps after getting rid of weak types altogether.
  jaxpr, out_avals, consts = _create_jaxpr(init_avals)
  if config.mutable_array_checks.value:
    _check_no_aliased_closed_over_refs(dbg_body, consts, list(args))
  carry_out_avals, ys_avals = out_avals.unpack()
  if len(carry_out_avals) != len(init_avals):
    _check_carry_type('scan body', f, init_avals, carry_out_avals)
  init_flat, changed = init_flat.map3(
     init_avals, carry_out_avals,

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Change the body to return exactly (carry, y)
2. If no per-step output, return (carry, None)
3. Verify both branches of any internal conditional return a 2-element pytree

Example fix

// before
def body(c, x):
  return c + x
carry, ys = lax.scan(body, 0, xs)
// after
def body(c, x):
  return c + x, x * 2
carry, ys = lax.scan(body, 0, xs)

out = body(carry_example, x_example)
assert isinstance(out, tuple) and len(out) == 2, 'scan body must return (carry, y)'

def returns_pair(body, c_ex, x_ex) -> bool:
    try:
        c, y = body(c_ex, x_ex)
        return True
    except (TypeError, ValueError):
        return False

try: lax.scan(body, init, xs)
except TypeError as e:
    if 'must be a pair' in str(e): fix body to return (carry, y) and rerun
    else: raise

When it happens

Trigger: def body(c, x): return c + x (single value), or returning (c, y, extra), or returning a dict; lax.scan raises during jaxpr creation.

Common situations: First-time scan users adapting a for-loop body; refactoring code where the body used to return only the carry; inconsistent returns between code paths in the body.

Related errors

• The inputs to the closure produced by closure_convert must h
• lax.scan: f argument should be a callable.
• scan got value with no leading axis to scan over: {}.
• {} function carry input and carry output must have the same
• numpy masked arrays are not supported as direct inputs to JA

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