jax-ml/jax · error · TypeError

{} function carry input and carry output must have equal typ

Error message

{} function carry input and carry output must have equal types, but they differ:

{}
{}Revise the function so that all output types match the corresponding input types.

What it means

Even when the carry pytree structure matches, JAX requires each output leaf of a scan/while_loop body to have the same avals (shape and dtype) as the corresponding input leaf, since loop state must be typed identically every iteration. This error lists the per-leaf differences (shape/dtype mismatches) and, under shard_map, may note varying manual axes (VMA) mismatches between input and output carries.

Source

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

        for path, in_aval, out_aval in zip(in_carry.paths, in_carry, out_carry)
        if not core.typematch(in_aval, out_aval) and
        isinstance(in_aval, ShapedArray) and isinstance(out_aval, ShapedArray)
        and in_aval.mat.varying != out_aval.mat.varying
        and out_aval.mat.varying - in_aval.mat.varying]

    if not pvary_applications:
      pvary_msg = ''
    elif len(pvary_applications) == 1:
      pvary_msg = f'This might be fixed by {pvary_applications[0]}.\n'
    else:
      pvary_msg = ('This might be fixed by:\n' +
                   '\n'.join(f'  * {d};\n' for d in pvary_applications[:-1])
                   + f'  * {pvary_applications[-1]}.\n')
    if pvary_msg:
      pvary_msg += ("See https://docs.jax.dev/en/latest/notebooks/shard_map.html#scan-vma "
                    "for more information.\n\n")

    raise TypeError(
        f"{name} function carry input and carry output must have equal types, "
        "but they differ:\n\n"
        f"{differences}\n"
        f"{pvary_msg}"
        "Revise the function so that all output types match the corresponding "
        "input types.")

# TODO(mattjj): re-land #19819 version? simpler, but caused ~1 perf regression.
def _scan_impl(*args, reverse, length, ft_in, ft_out, jaxpr,
               unroll):
  consts, carry, xs_ = _map(list, ft_in.update(args).unpack())
  _, y_avals = ft_out.update(jaxpr.out_avals).unpack()
  if unroll == 0:
    num_trips, remainder = 0, length
  else:
    num_trips, remainder = divmod(length, unroll)

  xs_rem: tuple[Array, ...] = ()

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make each carry output leaf match input dtype/shape: cast explicitly with jax.lax.convert_element_type(out, carry.dtype) or .astype before returning
Check the printed per-leaf diffs to find which leaf index mismatches and fix that computation (e.g. pass dtype= to jnp.sum/jnp.zeros)
If shapes differ, reshape the output back to the input's shape inside body_fun
For the VMA variant under shard_map, ensure manual axes declared for scan inputs match outputs, or as a temporary workaround pass check_vma=False to jax.shard_map

Example fix

// before
def body(c, x):
    return c + x.sum()  # sum upcasts bf16 -> f32
jax.lax.scan(body, jnp.zeros((), jnp.bfloat16), xs)

// after
def body(c, x):
    return c + x.sum(dtype=jnp.bfloat16)  # or lax.convert_element_type(c + x.sum(), c.dtype)
jax.lax.scan(body, jnp.zeros((), jnp.bfloat16), xs)

Defensive patterns

Strategy: validation

Validate before calling

import jax, numpy as np
init_avals = jax.api_util.flatten_axes  # simpler: trace once
jaxpr = jax.make_jaxpr(lambda c: body_fun(c))(init_val)
in_l = jax.tree_util.tree_leaves(init_val)
out_l = jax.tree_util.tree_leaves(body_fun(init_val))
for a, b in zip(in_l, out_l):
    assert jnp.shape(a) == jnp.shape(b) and jnp.dtype(a) == jnp.dtype(b)

Try / catch

try:
    jax.lax.while_loop(cond, body, init)
except TypeError as e:
    if 'equal types' in str(e):
        # add explicit casts in body and retry
        raise

Prevention

Set explicit dtype= on reductions and initializers inside loop bodies
End body_fun with an explicit cast of the carry: lax.convert_element_type(out, init.dtype) when in mixed precision
Lock in tree_structure plus leaf dtype/shape checks in unit tests for loop bodies

When it happens

Trigger: body_fun promotes precision (e.g. jnp.sum with default dtype, or x / y widening to float32) so the returned carry has dtype float32 while init is float16/bfloat16; operations that change shape (reshape, squeeze on the carry); under jax.shard_map, scan bodies whose input and output carries have inconsistent varying manual axes when check_vma=True.

Common situations: Mixed-precision training where init carry is bf16 but an op in the body upcasts to f32; accumulating with jnp.zeros_like on a different dtype; forgetting dtype= on reductions; shard_map + scan compositions after a JAX upgrade tightening VMA checks.

Related errors

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