{"record":{"id":"029114c381d65c8c","repo":"jax-ml/jax","slug":"scan-carry-input-and-output-got-mismatched-varying","errorCode":null,"errorMessage":"Scan carry input and output got mismatched varying manual axes {} and {}. Please open an issue at https://github.com/jax-ml/jax/issues, and as a temporary workaround pass the check_vma=False argument to `jax.shard_map`","messagePattern":"Scan carry input and output got mismatched varying manual axes (.+?) and (.+?)\\. Please open an issue at https://github\\.com/jax-ml/jax/issues, and as a temporary workaround pass the check_vma=False argument to `jax\\.shard_map`","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/lax/control_flow/loops.py","lineNumber":741,"sourceCode":"  #                   memory_space=aval.memory_space)\n  # return core.pvary(empty, tuple(aval.mat.varying))\n  empty = core.pvary(lax.empty2(aval.dtype, memory_space=aval.memory_space),\n                     tuple(aval.mat.varying))\n  with use_abstract_mesh(sharding.mesh):\n    out = lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)\n  return out\n\n\ndef _scan_abstract_eval(*args, reverse, length, ft_in, ft_out, jaxpr,\n                        unroll):\n  if len(args) != len(jaxpr.in_avals):\n    raise ValueError(\"scan number of arguments doesn't match the number \"\n                     \"of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}\")\n  out_carry_avals, y_avals = ft_out.update(jaxpr.out_avals).unpack()\n  _, in_carry_avals, _ = ft_in.update(args).unpack()\n  if ([i.mat for i in in_carry_avals if isinstance(i, core.ShapedArray)] !=\n      [o.mat for o in out_carry_avals if isinstance(o, core.ShapedArray)]):\n    raise ValueError(\n        'Scan carry input and output got mismatched varying manual axes '\n        f'{in_carry_avals} and {out_carry_avals}. Please open an '\n        'issue at https://github.com/jax-ml/jax/issues, and as a '\n        'temporary workaround pass the check_vma=False argument to '\n        '`jax.shard_map`')\n  ys_avals = _map(partial(core.unmapped_leading_aval, length), y_avals)\n  return list(out_carry_avals) + list(ys_avals), core.positional_effects(jaxpr)\n\ndef _scan_jvp(primals, tangents, reverse, length, jaxpr, ft_in, ft_out, unroll):\n  nonzeros = [type(t) is not ad_util.Zero for t in tangents]\n  const_nz, init_nz, xs_nz = ft_in.update(nonzeros).unpack()\n\n  # Fixpoint computation of which carry are not ad.zero: either\n  # non-zero from init, or the carry out is non-zero. Each iteration promotes\n  # at least one carry to non-zero. We need at most len(carry) iterations,\n  # but we need one last iteration to prepare the jaxpr based on the final\n  # carry_nz.\n  carry_nz = init_nz","sourceCodeStart":723,"sourceCodeEnd":759,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/lax/control_flow/loops.py#L723-L759","documentation":"Under jax.shard_map, scan carries can be annotated with varying manual axes (VMA) describing how they vary across shards; JAX verifies that a carry's input and output VMA specifications are consistent. This error means the input and output carry avals have mismatched varying-axes matrices, which the JAX authors consider an unsupported/inconsistent state — the message explicitly asks you to file a bug and offers check_vma=False as a workaround.","triggerScenarios":"Using jax.lax.scan inside jax.shard_map(check_vma=True) where operations in the body (e.g. ppermute, all-gather-like reshapes, or dtype/aval-affecting ops) change the varying-axes of the carry between input and output; combos of shard_map with remat, vmap, or custom collectives on the carry.","commonSituations":"Multi-host/multi-device sharded training loops after upgrading JAX (VMA checking was tightened in newer releases); writing custom collectives inside sharded scan bodies; mismatched mesh axis annotations on input vs output state.","solutions":["Pass check_vma=False to jax.shard_map as the documented temporary workaround and verify numerics manually","Try to make the body preserve the carry's sharding/varying axes (avoid ops that change sharding of the carry without resharding it back)","Minimize the example and open an issue at https://github.com/jax-ml/jax/issues as the message requests","Upgrade to the latest JAX — several VMA false-positive bugs were fixed across releases"],"exampleFix":"// before\njax.shard_map(f, mesh, in_specs=..., out_specs=...)(x)  # f contains lax.scan\n\n// after\njax.shard_map(f, mesh, in_specs=..., out_specs=..., check_vma=False)(x)","handlingStrategy":"fallback","validationCode":"out = jax.shard_map(f, mesh, in_specs=..., out_specs=..., check_vma=False)(x)  # use when hitting VMA false positives","typeGuard":null,"tryCatchPattern":"try:\n    out = jax.shard_map(f, mesh, in_specs=spec, out_specs=ospec)(x)\nexcept ValueError as e:\n    if 'varying manual axes' in str(e):\n        out = jax.shard_map(f, mesh, in_specs=spec, out_specs=ospec, check_vma=False)(x)\n        # verify numerics manually against a single-device run","preventionTips":["Keep carry sharding annotations identical at scan input and output inside shard_map","Reproduce sharded numerics single-device first so check_vma=False fallback is safe","Watch JAX release notes for shard_map/VMA changes before upgrading"],"tags":["jax","shard-map","vma","scan","spmd","multi-device"],"backgroundTag":"jax-shard-map-vma-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}