{"record":{"id":"8d0a447c51d3e555","repo":"jax-ml/jax","slug":"scan-body-output-must-be-a-pair-got","errorCode":null,"errorMessage":"scan body output must be a pair, got {}.","messagePattern":"scan body output must be a pair, got (.+?)\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":384,"sourceCode":"    for i in maybe_reversed(range(length)):\n      xs_slice = xs_flat.map(lambda x: slicing.index_in_dim(x, i, keepdims=False))\n      carry, y = f(carry, xs_slice.unflatten())\n      ys.append(y)\n    stack = lambda *ys: _stack(ys)\n    stacked_y = tree_map(stack, *maybe_reversed(ys))\n    return carry, stacked_y\n\n  if config.mutable_array_checks.value:\n    check_no_aliased_ref_args(lambda: dbg_body, list(args_avals), list(args))\n\n  x_avals = xs_avals.map(lambda aval: core.mapped_leading_aval(length, aval))\n  def _create_jaxpr(carry_avals):\n    new_arg_avals = ft.pack(((carry_avals, x_avals), {}))\n    jaxpr, out_avals = pe.trace_to_jaxpr(f, new_arg_avals, dbg_body)\n    jaxpr, consts = pe.separate_consts(jaxpr)\n    if not out_avals.unpackable or len(out_avals.unpack()) != 2:\n      msg = \"scan body output must be a pair, got {}.\"\n      raise TypeError(msg.format(out_avals.unflatten()))\n    return jaxpr, out_avals, consts\n\n  # The carry input and output avals must match exactly. However, we want to account for\n  # the case when init contains weakly-typed values (e.g. Python scalars), with avals that\n  # may not match the output despite being compatible by virtue of their weak type.\n  # To do this, we compute the jaxpr in two passes: first with the raw inputs, and if\n  # necessary, a second time with modified init values.\n  # TODO(dougalm): this two-pass stuff is expensive (exponential in scan nesting\n  # depth) and incomplete (because in the general case it takes more than two passes).\n  # Let's get rid of it, perhaps after getting rid of weak types altogether.\n  jaxpr, out_avals, consts = _create_jaxpr(init_avals)\n  if config.mutable_array_checks.value:\n    _check_no_aliased_closed_over_refs(dbg_body, consts, list(args))\n  carry_out_avals, ys_avals = out_avals.unpack()\n  if len(carry_out_avals) != len(init_avals):\n    _check_carry_type('scan body', f, init_avals, carry_out_avals)\n  init_flat, changed = init_flat.map3(\n     init_avals, carry_out_avals,","sourceCodeStart":366,"sourceCodeEnd":402,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L366-L402","documentation":"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.","triggerScenarios":"def body(c, x): return c + x (single value), or returning (c, y, extra), or returning a dict; lax.scan raises during jaxpr creation.","commonSituations":"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.","solutions":["Change the body to return exactly (carry, y)","If no per-step output, return (carry, None)","Verify both branches of any internal conditional return a 2-element pytree"],"exampleFix":"// before\ndef body(c, x):\n  return c + x\ncarry, ys = lax.scan(body, 0, xs)\n// after\ndef body(c, x):\n  return c + x, x * 2\ncarry, ys = lax.scan(body, 0, xs)","handlingStrategy":"validation","validationCode":"out = body(carry_example, x_example)\nassert isinstance(out, tuple) and len(out) == 2, 'scan body must return (carry, y)'","typeGuard":"def returns_pair(body, c_ex, x_ex) -> bool:\n    try:\n        c, y = body(c_ex, x_ex)\n        return True\n    except (TypeError, ValueError):\n        return False","tryCatchPattern":"try: lax.scan(body, init, xs)\nexcept TypeError as e:\n    if 'must be a pair' in str(e): fix body to return (carry, y) and rerun\n    else: raise","preventionTips":["Always write scan bodies as 'return carry, y'","Return (carry, None) when there is no per-step output","Unit-test bodies with example inputs before composing into scans"],"tags":["jax","scan","typeerror","return-shape","pytree"],"backgroundTag":"function-return-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}