jax-ml/jax · error · ValueError

Scan carry input and output got mismatched varying manual ax

Error message

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`

What it means

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.

Source

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

  #                   memory_space=aval.memory_space)
  # return core.pvary(empty, tuple(aval.mat.varying))
  empty = core.pvary(lax.empty2(aval.dtype, memory_space=aval.memory_space),
                     tuple(aval.mat.varying))
  with use_abstract_mesh(sharding.mesh):
    out = lax.broadcast(empty, (*prefix, *aval.shape), out_sharding=sharding)
  return out


def _scan_abstract_eval(*args, reverse, length, ft_in, ft_out, jaxpr,
                        unroll):
  if len(args) != len(jaxpr.in_avals):
    raise ValueError("scan number of arguments doesn't match the number "
                     "of jaxpr arguments: {len(args)} vs {len(jaxpr.in_avals)}")
  out_carry_avals, y_avals = ft_out.update(jaxpr.out_avals).unpack()
  _, in_carry_avals, _ = ft_in.update(args).unpack()
  if ([i.mat for i in in_carry_avals if isinstance(i, core.ShapedArray)] !=
      [o.mat for o in out_carry_avals if isinstance(o, core.ShapedArray)]):
    raise ValueError(
        'Scan carry input and output got mismatched varying manual axes '
        f'{in_carry_avals} and {out_carry_avals}. 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`')
  ys_avals = _map(partial(core.unmapped_leading_aval, length), y_avals)
  return list(out_carry_avals) + list(ys_avals), core.positional_effects(jaxpr)

def _scan_jvp(primals, tangents, reverse, length, jaxpr, ft_in, ft_out, unroll):
  nonzeros = [type(t) is not ad_util.Zero for t in tangents]
  const_nz, init_nz, xs_nz = ft_in.update(nonzeros).unpack()

  # Fixpoint computation of which carry are not ad.zero: either
  # non-zero from init, or the carry out is non-zero. Each iteration promotes
  # at least one carry to non-zero. We need at most len(carry) iterations,
  # but we need one last iteration to prepare the jaxpr based on the final
  # carry_nz.
  carry_nz = init_nz

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass check_vma=False to jax.shard_map as the documented temporary workaround and verify numerics manually
2. Try to make the body preserve the carry's sharding/varying axes (avoid ops that change sharding of the carry without resharding it back)
3. Minimize the example and open an issue at https://github.com/jax-ml/jax/issues as the message requests
4. Upgrade to the latest JAX — several VMA false-positive bugs were fixed across releases

Example fix

// before
jax.shard_map(f, mesh, in_specs=..., out_specs=...)(x)  # f contains lax.scan

// after
jax.shard_map(f, mesh, in_specs=..., out_specs=..., check_vma=False)(x)

out = jax.shard_map(f, mesh, in_specs=..., out_specs=..., check_vma=False)(x)  # use when hitting VMA false positives

try:
    out = jax.shard_map(f, mesh, in_specs=spec, out_specs=ospec)(x)
except ValueError as e:
    if 'varying manual axes' in str(e):
        out = jax.shard_map(f, mesh, in_specs=spec, out_specs=ospec, check_vma=False)(x)
        # verify numerics manually against a single-device run

When it happens

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

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

Related errors

• {} function carry input and carry output must have equal typ
• Primitive {prim_name} requires varying manual axes to match,
• 0th dimension of all xs should be replicated. Got {}
• Explicit sharding inference for ragged_dot_general is not cu
• {name} cannot accept args with unreduced_kind={mat.unreduced

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