jax-ml/jax · error · ValueError

The context mesh should match the mesh passed to shard_map

Error message

The context mesh {ctx_mesh} should match the mesh passed to shard_map {mesh}

What it means

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.

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

Example fix

// before
with jax.set_mesh(mesh_a):
    shard_map(f, mesh=mesh_b, in_specs=P('i'), out_specs=P('i'))(x)

// after
shard_map(f, mesh=mesh_b, in_specs=P('i'), out_specs=P('i'))(x)  # outside the context
# or ensure mesh_b == mesh_a / omit mesh= inside the context

Defensive patterns

Strategy: validation

Validate before calling

ctx = jax.sharding.get_abstract_mesh()
if mesh is not None and not ctx.empty and mesh.abstract_mesh != ctx:
    raise ValueError('provided mesh differs from active context mesh; pick one')

Prevention

Use one mechanism per call site: explicit mesh OR context, not both
Name a single mesh-creation helper and reuse it everywhere

When it happens

Trigger: 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.

Common situations: 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.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/3a1692ca8b157d4a. Report an issue: GitHub.

Appendix: source

Thrown at jax/_src/shard_map.py:374


def _axes_to_pspec(axis_name, axis):
  if axis is None:
    return P()
  return P(*[None] * axis + [axis_name])


def _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):
  if mesh is None:
    mesh = get_abstract_mesh()
    if mesh.empty:
      raise ValueError(
          "The context mesh cannot be empty. Use"
          " `jax.set_mesh(mesh)` to enter into a mesh context")
  else:
    ctx_mesh = get_abstract_mesh()
    if not ctx_mesh.empty and mesh.abstract_mesh != ctx_mesh:
      raise ValueError(
          f"The context mesh {ctx_mesh} should match the mesh passed to"
          f" shard_map {mesh}")

  if not isinstance(mesh, (Mesh, AbstractMesh)):
    raise TypeError("shard_map requires a `jax.sharding.Mesh` or a "
                    "`jax.sharding.AbstractMesh` instance for its "
                    f"second argument, but got {mesh} of type {type(mesh)}.")
  if mesh.empty:
    raise ValueError(f"shard_map requires a non-empty mesh. Got {mesh}")

  mesh_axis_names_wo_vmap = (
      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names
  )

  if not isinstance(axis_names, (frozenset, set)):
    raise TypeError(
        "`axis_names` argument of shard_map should be of type `frozenset` or"
        f" `set`. Got type: {type(axis_names)}")

View on GitHub (pinned to 1e1c6a8fc0)