{"record":{"id":"3a1692ca8b157d4a","repo":"jax-ml/jax","slug":"the-context-mesh-ctx-mesh-should-match-the-mesh","errorCode":null,"errorMessage":"The context mesh {ctx_mesh} should match the mesh passed to shard_map {mesh}","messagePattern":"The context mesh (.+?) should match the mesh passed to shard_map (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":374,"sourceCode":"\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:\n    ctx_mesh = get_abstract_mesh()\n    if not ctx_mesh.empty and mesh.abstract_mesh != ctx_mesh:\n      raise ValueError(\n          f\"The context mesh {ctx_mesh} should match the mesh passed to\"\n          f\" shard_map {mesh}\")\n\n  if not isinstance(mesh, (Mesh, AbstractMesh)):\n    raise TypeError(\"shard_map requires a `jax.sharding.Mesh` or a \"\n                    \"`jax.sharding.AbstractMesh` instance for its \"\n                    f\"second argument, but got {mesh} of type {type(mesh)}.\")\n  if mesh.empty:\n    raise ValueError(f\"shard_map requires a non-empty mesh. Got {mesh}\")\n\n  mesh_axis_names_wo_vmap = (\n      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names\n  )\n\n  if not isinstance(axis_names, (frozenset, set)):\n    raise TypeError(\n        \"`axis_names` argument of shard_map should be of type `frozenset` or\"\n        f\" `set`. Got type: {type(axis_names)}\")","sourceCodeStart":356,"sourceCodeEnd":392,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L356-L392","documentation":"If you pass an explicit mesh to shard_map while a non-empty context mesh (from jax.set_mesh) is active, JAX requires the two to match; otherwise the mapping of partition specs to devices would be ambiguous. This ValueError reports the mismatch between the context mesh and the provided mesh.","triggerScenarios":"Entering `with jax.set_mesh(mesh_a):` then calling shard_map(f, mesh=mesh_b, ...) where mesh_b's abstract layout differs from mesh_a — different axis names, sizes, or types.","commonSituations":"Refactoring code that mixes the old explicit-mesh API with the new context-mesh API; nested contexts from library code (e.g. a training loop sets a mesh while a utility passes its own); reusing shard_map-wrapped functions across different device meshes.","solutions":["Make both meshes identical (same abstract mesh), or drop one: either rely on the context and omit mesh=, or exit the context before passing mesh explicitly","Audit jax.set_mesh usage in library/framework layers wrapping your call"],"exampleFix":"// before\nwith jax.set_mesh(mesh_a):\n    shard_map(f, mesh=mesh_b, in_specs=P('i'), out_specs=P('i'))(x)\n\n// after\nshard_map(f, mesh=mesh_b, in_specs=P('i'), out_specs=P('i'))(x)  # outside the context\n# or ensure mesh_b == mesh_a / omit mesh= inside the context","handlingStrategy":"validation","validationCode":"ctx = jax.sharding.get_abstract_mesh()\nif mesh is not None and not ctx.empty and mesh.abstract_mesh != ctx:\n    raise ValueError('provided mesh differs from active context mesh; pick one')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Use one mechanism per call site: explicit mesh OR context, not both","Name a single mesh-creation helper and reuse it everywhere"],"tags":["jax","shard-map","mesh","context-mismatch","distributed"],"backgroundTag":"mesh-context-conflict","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}