jax-ml/jax · error · ValueError

Primitive {prim_name} requires varying manual axes to match,

Error message

Primitive {prim_name} requires varying manual axes to match, but got {[vma, *vmas]}. 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

SPMD primitives require all operand avals to share the same 'varying manual axes' (vma). When operands disagree, JAX raises this internal invariant error and points to the check_vma=False escape hatch on jax.shard_map.

Source

Thrown at jax/_src/core.py:2820

      # Will need more changes to pvary to allow such partialness.
      if src_reduced == rest_vma:
        out.append(
            reduced_vary_cast(arg, tuple(n for n in out_vma if n in rest_vma)))
      else:
        out.append(pvary(arg, tuple(n for n in out_vma if n in rest_vma)))
    else:
      out.append(arg)
  return out

def standard_vma_rule(prim_name, *avals, **kwargs) -> frozenset[AxisName]:
  if not config._check_vma.value:
    return frozenset()
  avals = tuple(a for a in avals if a is not abstract_token)
  if not avals:
    return frozenset()
  vma, *vmas = (a.mat.varying for a in avals)
  if not all(vma == vma_ for vma_ in vmas):
    raise ValueError(
        f'Primitive {prim_name} requires varying manual axes '
        f'to match, but got {[vma, *vmas]}. 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`')
  return vma

@dataclass(frozen=True, slots=True)
class bint(dtypes.ExtendedDType):
  bound: int

  @property
  def type(self) -> type:
    return dtypes.extended

  @property
  def name(self) -> str:
    return f'bint{{≤{self.bound}}}'

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. As the message says: pass check_vma=False to jax.shard_map as a temporary workaround and file the issue
2. Upgrade/downgrade JAX — this is usually fixed quickly
3. Simplify the shard_map inputs so all operands share the same manual axes

Example fix

// before
jax.shard_map(fn, mesh, in_specs=(P('x', None), P(None, 'x')))(a, b)

// after
jax.shard_map(fn, mesh, in_specs=(P('x', None), P(None, 'x')), check_vma=False)(a, b)

try:
    out = jax.shard_map(fn, mesh, in_specs)(a, b)
except ValueError as e:
    if 'check_vma=False' in str(e):
        out = jax.shard_map(fn, mesh, in_specs, check_vma=False)(a, b)
    else:
        raise

When it happens

Trigger: A multi-operand primitive (e.g. an elementwise op inside shard_map) receives operands whose mat.varying differ — typically due to a JAX bug or unusual manual_axis_type plumbing inside shard_map.

Common situations: Hitting a genuine JAX spmd bug after version upgrades; using shard_map with mixed shardings/manual axes; custom primitives with mismatched operand avals.

Related errors

• Scan carry input and output got mismatched varying manual ax
• {name} cannot accept args with unreduced_kind={mat.unreduced
• vmap spmd_axis_name cannot appear in shard_map in_specs
• vmap spmd_axis_name cannot appear in shard_map out_specs
• callbacks are only supported in spmd computations when all m

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