{"record":{"id":"de2dbe97819a0bee","repo":"jax-ml/jax","slug":"primitive-prim-name-requires-varying-manual-axes","errorCode":null,"errorMessage":"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`","messagePattern":"Primitive (.+?) requires varying manual axes to match, but got (.+?)\\. 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/core.py","lineNumber":2820,"sourceCode":"      # Will need more changes to pvary to allow such partialness.\n      if src_reduced == rest_vma:\n        out.append(\n            reduced_vary_cast(arg, tuple(n for n in out_vma if n in rest_vma)))\n      else:\n        out.append(pvary(arg, tuple(n for n in out_vma if n in rest_vma)))\n    else:\n      out.append(arg)\n  return out\n\ndef standard_vma_rule(prim_name, *avals, **kwargs) -> frozenset[AxisName]:\n  if not config._check_vma.value:\n    return frozenset()\n  avals = tuple(a for a in avals if a is not abstract_token)\n  if not avals:\n    return frozenset()\n  vma, *vmas = (a.mat.varying for a in avals)\n  if not all(vma == vma_ for vma_ in vmas):\n    raise ValueError(\n        f'Primitive {prim_name} requires varying manual axes '\n        f'to match, but got {[vma, *vmas]}. Please open an issue at '\n        'https://github.com/jax-ml/jax/issues and as a temporary '\n        'workaround pass the check_vma=False argument to `jax.shard_map`')\n  return vma\n\n@dataclass(frozen=True, slots=True)\nclass bint(dtypes.ExtendedDType):\n  bound: int\n\n  @property\n  def type(self) -> type:\n    return dtypes.extended\n\n  @property\n  def name(self) -> str:\n    return f'bint{{≤{self.bound}}}'\n","sourceCodeStart":2802,"sourceCodeEnd":2838,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/core.py#L2802-L2838","documentation":"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.","triggerScenarios":"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.","commonSituations":"Hitting a genuine JAX spmd bug after version upgrades; using shard_map with mixed shardings/manual axes; custom primitives with mismatched operand avals.","solutions":["As the message says: pass check_vma=False to jax.shard_map as a temporary workaround and file the issue","Upgrade/downgrade JAX — this is usually fixed quickly","Simplify the shard_map inputs so all operands share the same manual axes"],"exampleFix":"// before\njax.shard_map(fn, mesh, in_specs=(P('x', None), P(None, 'x')))(a, b)\n\n// after\njax.shard_map(fn, mesh, in_specs=(P('x', None), P(None, 'x')), check_vma=False)(a, b)","handlingStrategy":"fallback","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    out = jax.shard_map(fn, mesh, in_specs)(a, b)\nexcept ValueError as e:\n    if 'check_vma=False' in str(e):\n        out = jax.shard_map(fn, mesh, in_specs, check_vma=False)(a, b)\n    else:\n        raise","preventionTips":["Keep operand manual axes consistent in shard_map specs","Pin a known-good JAX version; report the bug upstream"],"tags":["jax","shard-map","spmd","internal-error","workaround"],"backgroundTag":"spmd-axis-misuse","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}