jax-ml/jax · error · TypeError

{} function carry input and carry output must have the same

Error message

{} function carry input and carry output must have the same pytree structure, but they differ:

{}
Revise the function so that the carry output has the same pytree structure as the carry input.

What it means

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.

Source

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

    except:
      out_carry_unflat = None

    if out_carry_unflat is None:
      differences = (f'the input tree structure is:\n{in_carry.tree}\n' +
                     f'the output tree structure is:\n{out_carry.tree}\n')
    else:
      diffs = [f'{component(path)} is a {thing1} but the corresponding component '
               f'of the carry output is a {thing2}, so {explanation}'
               for path, thing1, thing2, explanation
               in equality_errors(in_carry.unflatten(), out_carry.unflatten())]
      if len(diffs) == 0:
        return  # the trees may have different aux data, but structures are same
      elif len(diffs) == 1:
        differences = f'{_capitalize(diffs[0])}.\n'
      else:
        differences = ('\n'.join(f'  * {d};\n' for d in diffs[:-1])
                       + f'  * {diffs[-1]}.\n')
    raise TypeError(
        f"{name} function carry input and carry output must have the same "
        "pytree structure, but they differ:\n\n"
        f"{differences}\n"
        "Revise the function so that the carry output has the same pytree "
        "structure as the carry input.")
  if not all(_map(core.typematch, in_carry, out_carry)):
    diffs = [f'{component(path)} has type {in_aval.str_short()}'
             ' but the corresponding output carry component has type '
             f'{out_aval.str_short()}'
             f'{core.aval_mismatch_extra(in_aval, out_aval)}'
             for path, in_aval, out_aval in zip(in_carry.paths, in_carry, out_carry)
             if not core.typematch(in_aval, out_aval)]

    if len(diffs) == 0:
      return  # seems unreachable but in any case we don't have a good error msg
    if len(diffs) == 1:
      differences = f'{_capitalize(diffs[0])}.\n'
    else:

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

// before
carry = 0.0
def body(c, x):
    return (c + x, x)  # returns tuple, init is scalar
jax.lax.scan(body, carry, xs)

// after
carry = (0.0, None)  # match structure, or simplify body
def body(c, x):
    return c + x, x  # carry stays scalar
jax.lax.scan(body, carry, xs)

Defensive patterns

Strategy: validation

Validate before calling

import jax
struct_in = jax.tree_util.tree_structure(init_val)
struct_out = jax.tree_util.tree_structure(body_fun(init_val))  # or f(init_val, xs[0])[0] for scan
assert struct_in == struct_out, f'{struct_in} vs {struct_out}'

Type guard

def carry_structures_match(init, body, sample_x=None) -> bool:
    out = body(init) if sample_x is None else body(init, sample_x)[0]
    return jax.tree_util.tree_structure(init) == jax.tree_util.tree_structure(out)

Try / catch

try:
    result = jax.lax.scan(f, init, xs)
except TypeError as e:
    if 'pytree structure' in str(e):
        # inspect trees and fix body return structure
        raise

Prevention

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

When it happens

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

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

Related errors

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