jax-ml/jax · error · ValueError

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

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.

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)

Solutions

1. 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
2. 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

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')

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

• The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t
• shard_map requires a `jax.sharding.Mesh` or a `jax.sharding.
• shard_map requires a non-empty mesh. Got {mesh}
• jax.shard_map requires axis_names={axis_names} to be a subse
• Mapped away dimension of inputs passed to vmap should be sha

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