jax-ml/jax · error · ValueError

The context mesh cannot be empty. Use `jax.set_mesh(mesh)` t

Error message

The context mesh cannot be empty. Use `jax.set_mesh(mesh)` to enter into a mesh context

What it means

When shard_map/smap is called without an explicit mesh argument, JAX falls back to the context (abstract) mesh via jax.set_mesh. If that context mesh is empty — no mesh context manager active — this ValueError is raised telling you to enter one with jax.set_mesh(mesh).

Source

Thrown at jax/_src/shard_map.py:368

    except _RepError as e:
      fails, out_tree, = e.args
      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)
      raise ValueError(msg) from None
    return out_ft.unflatten()
  return cast(F, wrapped)


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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Pass the mesh explicitly: shard_map(f, mesh=mesh, ...)
Wrap the call site in `with jax.set_mesh(mesh):` (and ensure jit-traced calls capture the context at trace time, not run time)

Example fix

// before
out = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))(x)  # no mesh, no context

// after
out = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)
# or
with jax.set_mesh(mesh):
    out = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))(x)

Defensive patterns

Strategy: validation

Validate before calling

from jax.sharding import AbstractMesh

def get_active_mesh(explicit=None):
    m = explicit
    if m is None:
        m = jax.sharding.get_abstract_mesh()
        if m.empty:
            raise ValueError('no context mesh; pass mesh= or use jax.set_mesh')
    return m

Prevention

Pass mesh= explicitly in library code
Scope `with jax.set_mesh(mesh):` around both tracing and execution
Check get_abstract_mesh().empty in test setup

When it happens

Trigger: Calling shard_map(f, in_specs=..., out_specs=...) with mesh=None outside of any `with jax.set_mesh(mesh):` block, or when Mesh context was exited before the deferred computation ran.

Common situations: Adopting the newer context-mesh API (replacing explicit mesh=), running under a SPMD multi-host setup where the context wasn't propagated, or calling shard_map lazily (inside jit) after the with-block expired.

Related errors

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