{"record":{"id":"beae67a657b1a051","repo":"jax-ml/jax","slug":"mesh-shape-of-the-input-a-sharding-mesh-shape-tup","errorCode":null,"errorMessage":"Mesh shape of the input {a.sharding.mesh.shape_tuple} does not match the mesh shape passed to shard_map  {mesh.shape_tuple} for shape {aval.str_short()}","messagePattern":"Mesh shape of the input (.+?) does not match the mesh shape passed to shard_map  (.+?) for shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":1194,"sourceCode":"  return mlir.wrap_with_shard_to_full_op(ctx, sx, aval_out, shard_proto,\n                                         unspecified)\n\ndef _pspec_mhlo_attrs(spec, aval: core.AbstractValue) -> str:\n  if isinstance(aval, core.ShapedArray):\n    names = _spec_to_names(spec)\n    return str(map(names.get, range(aval.ndim)))\n  return ''\n\n# Eager evaluation\n\ndef get_mesh_from_args(args_flat, mesh):\n  for a in args_flat:\n    if (hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding)\n        and not a.sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]\n\n      if a.sharding.mesh.shape_tuple != mesh.shape_tuple:\n        aval = core.shaped_abstractify(a)\n        raise ValueError(\n            f\"Mesh shape of the input {a.sharding.mesh.shape_tuple} does not\"\n            \" match the mesh shape passed to shard_map \"\n            f\" {mesh.shape_tuple} for shape {aval.str_short()}\")\n      mesh = a.sharding.mesh\n  if isinstance(mesh, AbstractMesh):\n    raise ValueError(\n        \"Please pass `jax.Array`s with a `NamedSharding` as input to\"\n        \" `shard_map` when passing `AbstractMesh` to the mesh argument.\")\n  assert isinstance(mesh, Mesh)\n  return mesh\n\ndef _spec_to_vma(spec):\n  return frozenset(p for s in spec.partitions if s is not None\n                   for p in (s if isinstance(s, tuple) else (s,)))\n\ndef _mat_to_spec(mesh, mat):\n  return P(order_wrt_mesh(mesh, mat.varying), unreduced=mat.unreduced,\n           reduced=mat.reduced, unreduced_kind=mat.unreduced_kind)","sourceCodeStart":1176,"sourceCodeEnd":1212,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L1176-L1212","documentation":"When shard_map infers the mesh from its arguments, every input array with a NamedSharding must live on a mesh with the same shape tuple as the mesh passed to (or previously inferred for) shard_map. Different mesh shapes make the sharding ambiguous.","triggerScenarios":"Passing arrays placed on a 2x4 mesh to shard_map(mesh=8x1 mesh, ...); typically when arrays come from a different jax.make_mesh/Mesh than the one used in the call.","commonSituations":"Notebooks or libraries where arrays were created under an earlier/different mesh layout; changing mesh definition between training stages without resharding.","solutions":["Re-create/place input arrays on the same mesh (shape tuple) as the shard_map mesh","Pass the same Mesh object used to build the arrays' NamedShardings to shard_map","Reshard via jax.device_put(arr, NamedSharding(target_mesh, spec)) before the call"],"exampleFix":"// before\nx = jax.device_put(x, NamedSharding(mesh_a, P('x')))\ny = jax.jit(shard_map(f, mesh_b, ...))(x)  # mesh_b shape differs\n// after\nx = jax.device_put(x, NamedSharding(mesh_b, P('x')))\ny = jax.jit(shard_map(f, mesh_b, ...))(x)","handlingStrategy":"validation","validationCode":"def same_mesh_shape(args, mesh):\n    return all(not hasattr(a, 'sharding') or not isinstance(a.sharding, NamedSharding)\n               or a.sharding.mesh.shape_tuple == mesh.shape_tuple for a in args)","typeGuard":"def args_on_mesh(args, mesh) -> bool:\n    return all(getattr(getattr(a, 'sharding', None), 'mesh', mesh).shape_tuple == mesh.shape_tuple for a in args)","tryCatchPattern":"try: shard_map(...) except ValueError as e: if 'Mesh shape' in str(e): jax.device_put args onto mesh and retry; else: raise","preventionTips":["Create one canonical mesh object per process and use it everywhere","Re-put arrays on the current mesh after changing mesh layouts"],"tags":["shard-map","mesh-mismatch","named-sharding","jax"],"backgroundTag":"mesh-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}