{"record":{"id":"242579eea18b6300","repo":"jax-ml/jax","slug":"input-sharding-mesh-aval-sharding-mesh-should-be","errorCode":null,"errorMessage":"Input sharding mesh {aval.sharding.mesh} should be equal to out_sharding mesh {out_sh.mesh}","messagePattern":"Input sharding mesh (.+?) should be equal to out_sharding mesh (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":2207,"sourceCode":"  new_invals = [next(ref_vals_) if isinstance(a, AbstractRef) else None\n                for a in ctx.in_avals]\n  assert next(ref_vals_, None) is None\n  return new_invals, out_vals\n\ndef _repspec(aval):\n  return aval.nospec(empty_abstract_mesh, False, ())\n\n# ----------------------- top level collectives --------------------------------\n\ndef _top_level_ag(x, aval, out_sh_, multi_dim):\n  assert aval.sharding.mesh.are_all_axes_explicit, aval.sharding.mesh\n  out_sh = canonicalize_sharding(out_sh_, \"top_level_all_gather\")\n  if out_sh is None:\n    raise ValueError(\n        f'out_sharding passed to top_level_all_gather cannot be {out_sh_}. It'\n        ' should be a PartitionSpec or NamedSharding.')\n  if aval.sharding.mesh != out_sh.mesh:\n    raise ValueError(\n        f'Input sharding mesh {aval.sharding.mesh} should be equal to'\n        f' out_sharding mesh {out_sh.mesh}')\n\n  in_spec = aval.sharding.spec\n  out_spec = out_sh.spec._normalized_spec_for_aval(len(in_spec))\n  if config.remove_size_one_mesh_axis_from_type.value:\n    out_spec = remove_size_one_mesh_axis_from_spec(out_spec, out_sh.mesh)\n\n  def f_shmap(x):\n    # Maybe this can just be 1 AG where we gather in a new dim and then do\n    # AG(new_dim) -> reshape -> transpose -> reshape but it might be expensive.\n    count = 0\n    for axis, (i, o) in enumerate(zip(in_spec.partitions, out_spec.partitions)):\n      if i == o:\n        continue\n      if not multi_dim and count > 0:\n        raise ValueError(\n            \"multiple dimensions cannot be all_gathered since multi_dim=False\"","sourceCodeStart":2189,"sourceCodeEnd":2225,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L2189-L2225","documentation":"top_level_all_gather requires the mesh implied by the input array's sharding to equal the mesh of the out_sharding. Collectives are emitted on the input's mesh, so a mismatched out_sharding mesh (different devices or axis names) is rejected.","triggerScenarios":"Input array sharded on mesh A (e.g. via jit in_shardings) but out_sharding built from a different mesh B, or meshes with same shape but different axis names.","commonSituations":"Creating multiple Mesh objects in one process; refactoring context meshes so names differ; mixing jax.make_mesh results across modules.","solutions":["Build out_sharding from the same Mesh object used for the input's sharding (or the context mesh)","Ensure mesh axis names match exactly — meshes differing only in names still compare unequal","Shard the input on the same mesh you intend to gather on before calling"],"exampleFix":"# before\nout = top_level_all_gather(x, NamedSharding(other_mesh, P('data')))\n\n# after\nout = top_level_all_gather(x, NamedSharding(x_mesh, P('data')))  # x_mesh = mesh x is sharded on","handlingStrategy":"validation","validationCode":"in_mesh = x.sharding.mesh if hasattr(x, 'sharding') else get_abstract_mesh()\nassert in_mesh == out_sharding.mesh if isinstance(out_sharding, NamedSharding) else True, 'mesh mismatch'","typeGuard":"def meshes_match(x, out_sharding) -> bool:\n    return isinstance(out_sharding, NamedSharding) and x.sharding.mesh == out_sharding.mesh","tryCatchPattern":null,"preventionTips":["Create one Mesh per hardware config and thread it through explicitly","Remember mesh equality includes axis names, not just shape"],"tags":["jax","mesh","all-gather","shard-map"],"backgroundTag":"jax-mesh-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}