{"record":{"id":"ffa9f92b75c4dc3b","repo":"jax-ml/jax","slug":"function-carry-input-and-carry-output-must-have","errorCode":null,"errorMessage":"{} function carry input and carry output must have the same pytree structure, but they differ:\n\n{}\nRevise the function so that the carry output has the same pytree structure as the carry input.","messagePattern":"(.+?) function carry input and carry output must have the same pytree structure, but they differ:\n\n(.+?)\nRevise the function so that the carry output has the same pytree structure as the carry input\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":566,"sourceCode":"    except:\n      out_carry_unflat = None\n\n    if out_carry_unflat is None:\n      differences = (f'the input tree structure is:\\n{in_carry.tree}\\n' +\n                     f'the output tree structure is:\\n{out_carry.tree}\\n')\n    else:\n      diffs = [f'{component(path)} is a {thing1} but the corresponding component '\n               f'of the carry output is a {thing2}, so {explanation}'\n               for path, thing1, thing2, explanation\n               in equality_errors(in_carry.unflatten(), out_carry.unflatten())]\n      if len(diffs) == 0:\n        return  # the trees may have different aux data, but structures are same\n      elif len(diffs) == 1:\n        differences = f'{_capitalize(diffs[0])}.\\n'\n      else:\n        differences = ('\\n'.join(f'  * {d};\\n' for d in diffs[:-1])\n                       + f'  * {diffs[-1]}.\\n')\n    raise TypeError(\n        f\"{name} function carry input and carry output must have the same \"\n        \"pytree structure, but they differ:\\n\\n\"\n        f\"{differences}\\n\"\n        \"Revise the function so that the carry output has the same pytree \"\n        \"structure as the carry input.\")\n  if not all(_map(core.typematch, in_carry, out_carry)):\n    diffs = [f'{component(path)} has type {in_aval.str_short()}'\n             ' but the corresponding output carry component has type '\n             f'{out_aval.str_short()}'\n             f'{core.aval_mismatch_extra(in_aval, out_aval)}'\n             for path, in_aval, out_aval in zip(in_carry.paths, in_carry, out_carry)\n             if not core.typematch(in_aval, out_aval)]\n\n    if len(diffs) == 0:\n      return  # seems unreachable but in any case we don't have a good error msg\n    if len(diffs) == 1:\n      differences = f'{_capitalize(diffs[0])}.\\n'\n    else:","sourceCodeStart":548,"sourceCodeEnd":584,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L548-L584","documentation":"In lax.scan and lax.while_loop, the carry (loop state) is passed from iteration to iteration, so JAX requires the value returned by the body function to have exactly the same pytree structure as the initial carry. This error is raised when the trees differ (different number of leaves, different nesting, dict keys, or tuple shapes). Because loop state cannot change shape between iterations in compiled code, JAX validates structure up front and fails fast.","triggerScenarios":"Calling jax.lax.scan(f, init, xs) or jax.lax.while_loop(cond, body, init_val) where body(init) returns a different pytree than init: e.g. init is a scalar but body returns a tuple, body returns (carry, y) from scan's f in the wrong order, body drops or adds a dict key, or a Python control-flow branch returns different structures.","commonSituations":"Accidentally swapping scan's expected return order (carry, ys); initializing carry as 0.0 but returning a tuple; refactoring state from a single array to a NamedTuple/dict without updating init; conditional accumulation inside while_loop bodies; returning jnp arrays vs Python scalars mixed across leaves.","solutions":["Restructure body_fun so its return value matches the pytree structure of the carry input exactly (same keys, same nesting, same leaf count)","For scan, double-check the function signature f(carry, x) -> (carry, y) and that you return them in that order","If the state legitimately changed shape, update the init_val to the new structure (e.g. wrap it in the same NamedTuple/dict)","Use jax.tree_util.tree_structure on init and on body(init) in a REPL/test to diff the two trees before calling the loop primitive"],"exampleFix":"// before\ncarry = 0.0\ndef body(c, x):\n    return (c + x, x)  # returns tuple, init is scalar\njax.lax.scan(body, carry, xs)\n\n// after\ncarry = (0.0, None)  # match structure, or simplify body\ndef body(c, x):\n    return c + x, x  # carry stays scalar\njax.lax.scan(body, carry, xs)","handlingStrategy":"validation","validationCode":"import jax\nstruct_in = jax.tree_util.tree_structure(init_val)\nstruct_out = jax.tree_util.tree_structure(body_fun(init_val))  # or f(init_val, xs[0])[0] for scan\nassert struct_in == struct_out, f'{struct_in} vs {struct_out}'","typeGuard":"def carry_structures_match(init, body, sample_x=None) -> bool:\n    out = body(init) if sample_x is None else body(init, sample_x)[0]\n    return jax.tree_util.tree_structure(init) == jax.tree_util.tree_structure(out)","tryCatchPattern":"try:\n    result = jax.lax.scan(f, init, xs)\nexcept TypeError as e:\n    if 'pytree structure' in str(e):\n        # inspect trees and fix body return structure\n        raise","preventionTips":["Define carry as a NamedTuple once and reuse it for init and returns","Add a unit test asserting tree_structure equality between init and body output before running full pipelines","Return explicit tuples in a fixed order; never build returns conditionally with different shapes"],"tags":["jax","pytree","scan","while-loop","carry-structure"],"backgroundTag":"pytree-structure-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}