{"record":{"id":"b3d0a73138ea751a","repo":"jax-ml/jax","slug":"shard-map-out-specs-vma-error-msg-from-inout-vma","errorCode":null,"errorMessage":"shard_map out_specs vma error (msg from _inout_vma_error)","messagePattern":"shard_map out_specs vma error \\(msg from _inout_vma_error\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":353,"sourceCode":"      return ans_ft.with_aux(out_specs_flat)\n\n    try:\n      newly_manual_axes = axis_names - set(mesh.manual_axes)\n      out_ft = shard_map_p.bind(\n          *dyn_args, subfuns=(f_wrapped,), mesh=mesh, in_specs=in_specs_flat,\n          check_vma=check_vma, newly_manual_axes=newly_manual_axes, debug_info=dbg)\n    except _SpecError as e:\n      fails, out_tree = e.args\n      msg = _spec_rank_error(SpecErrorType.out, f, out_tree, out_specs, fails)\n      if any(fail is not no_fail and not fail.shape for fail in fails):\n        msg += (\" In particular, for rank 0 outputs which are not constant \"\n                \"over the mesh, add at least one (singleton) axis to them so \"\n                \"that they can be concatenated using out_specs.\")\n      raise ValueError(msg) from None\n    except _RepError as e:\n      fails, out_tree, = e.args\n      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)\n      raise ValueError(msg) from None\n    return out_ft.unflatten()\n  return cast(F, wrapped)\n\n\ndef _axes_to_pspec(axis_name, axis):\n  if axis is None:\n    return P()\n  return P(*[None] * axis + [axis_name])\n\n\ndef _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):\n  if mesh is None:\n    mesh = get_abstract_mesh()\n    if mesh.empty:\n      raise ValueError(\n          \"The context mesh cannot be empty. Use\"\n          \" `jax.set_mesh(mesh)` to enter into a mesh context\")\n  else:","sourceCodeStart":335,"sourceCodeEnd":371,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L335-L371","documentation":"When shard_map's output reassembly encounters a value that is replicated across a mesh axis named in out_specs (i.e. the same data on every shard, so concatenation would duplicate it), a _RepError is raised internally and this ValueError with a message from _inout_vma_error is surfaced. It typically indicates a Virtual Mesh Axis (vmap-introduced axis) conflict or a replicated output being mapped over a named axis.","triggerScenarios":"A mapped function returns a value that does not vary along a mesh axis that out_specs names for it — e.g. broadcasting a constant across shards while out_specs=P('i'), often when using smap/vmap composition with shard_map.","commonSituations":"Mixing vmap/smap with shard_map where automatic spmd batching introduces implicit axes; returning mesh-broadcast constants; partially-updated JAX versions refining VMA error reporting.","solutions":["Restructure the function so the output genuinely varies along every mesh axis named in its out_specs entry, or drop that axis from the spec","Use jax.lax.axis_index / explicit per-shard values instead of broadcasting a constant","If the value is intentionally replicated, use out_specs that omit the replicated axis (or P() for full replication)"],"exampleFix":"// before\ndef f(x): return jnp.zeros_like(x)[0]  # same value on every shard\njax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)\n\n// after\ndef f(x): return x.sum(0)  # varies per shard is fine; or:\njax.shard_map(lambda x: jnp.zeros(1), mesh=mesh, in_specs=P('i'), out_specs=P())(x)","handlingStrategy":"validation","validationCode":"# ensure outputs vary along every mesh axis named in out_specs\ndef f(x):\n    per_shard = x.sum()          # constant across shards -> bad for P('i')\n    return per_shard + jnp.zeros(x.shape[0])  # shape tied to shard -> varies","typeGuard":null,"tryCatchPattern":"try:\n    out = shmapped(x)\nexcept ValueError:\n    out = jax.shard_map(lambda x: jnp.zeros(1), mesh=mesh,\n                        in_specs=P('i'), out_specs=P())(x)  # replicated fallback","preventionTips":["Never broadcast constants across a named out_specs axis","Compute outputs from shard-local data or lax.axis_index","Prefer P() for intentionally replicated outputs"],"tags":["jax","shard-map","out-specs","vma","replication"],"backgroundTag":"sharding-spec-rank-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}