{"record":{"id":"088934e661e4cb65","repo":"jax-ml/jax","slug":"function-carry-input-and-carry-output-must-have-088934","errorCode":null,"errorMessage":"{} function carry input and carry output must have equal types, but they differ:\n\n{}\n{}Revise the function so that all output types match the corresponding input types.","messagePattern":"(.+?) function carry input and carry output must have equal types, but they differ:\n\n(.+?)\n(.+?)Revise the function so that all output types match the corresponding input types\\.","errorType":"validation","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":611,"sourceCode":"        for path, in_aval, out_aval in zip(in_carry.paths, in_carry, out_carry)\n        if not core.typematch(in_aval, out_aval) and\n        isinstance(in_aval, ShapedArray) and isinstance(out_aval, ShapedArray)\n        and in_aval.mat.varying != out_aval.mat.varying\n        and out_aval.mat.varying - in_aval.mat.varying]\n\n    if not pvary_applications:\n      pvary_msg = ''\n    elif len(pvary_applications) == 1:\n      pvary_msg = f'This might be fixed by {pvary_applications[0]}.\\n'\n    else:\n      pvary_msg = ('This might be fixed by:\\n' +\n                   '\\n'.join(f'  * {d};\\n' for d in pvary_applications[:-1])\n                   + f'  * {pvary_applications[-1]}.\\n')\n    if pvary_msg:\n      pvary_msg += (\"See https://docs.jax.dev/en/latest/notebooks/shard_map.html#scan-vma \"\n                    \"for more information.\\n\\n\")\n\n    raise TypeError(\n        f\"{name} function carry input and carry output must have equal types, \"\n        \"but they differ:\\n\\n\"\n        f\"{differences}\\n\"\n        f\"{pvary_msg}\"\n        \"Revise the function so that all output types match the corresponding \"\n        \"input types.\")\n\n# TODO(mattjj): re-land #19819 version? simpler, but caused ~1 perf regression.\ndef _scan_impl(*args, reverse, length, ft_in, ft_out, jaxpr,\n               unroll):\n  consts, carry, xs_ = _map(list, ft_in.update(args).unpack())\n  _, y_avals = ft_out.update(jaxpr.out_avals).unpack()\n  if unroll == 0:\n    num_trips, remainder = 0, length\n  else:\n    num_trips, remainder = divmod(length, unroll)\n\n  xs_rem: tuple[Array, ...] = ()","sourceCodeStart":593,"sourceCodeEnd":629,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L593-L629","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\ndef body(c, x):\n    return c + x.sum()  # sum upcasts bf16 -> f32\njax.lax.scan(body, jnp.zeros((), jnp.bfloat16), xs)\n\n// after\ndef body(c, x):\n    return c + x.sum(dtype=jnp.bfloat16)  # or lax.convert_element_type(c + x.sum(), c.dtype)\njax.lax.scan(body, jnp.zeros((), jnp.bfloat16), xs)","handlingStrategy":"validation","validationCode":"import jax, numpy as np\ninit_avals = jax.api_util.flatten_axes  # simpler: trace once\njaxpr = jax.make_jaxpr(lambda c: body_fun(c))(init_val)\nin_l = jax.tree_util.tree_leaves(init_val)\nout_l = jax.tree_util.tree_leaves(body_fun(init_val))\nfor a, b in zip(in_l, out_l):\n    assert jnp.shape(a) == jnp.shape(b) and jnp.dtype(a) == jnp.dtype(b)","typeGuard":null,"tryCatchPattern":"try:\n    jax.lax.while_loop(cond, body, init)\nexcept TypeError as e:\n    if 'equal types' in str(e):\n        # add explicit casts in body and retry\n        raise","preventionTips":["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"],"tags":["jax","dtype-mismatch","scan","while-loop","shard-map","vma"],"backgroundTag":"dtype-shape-mismatch-in-loop-carry","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}