jax-ml/jax · error · ValueError

Context mesh cannot be empty. Please use `jax.set_mesh` API

Error message

Context mesh cannot be empty. Please use `jax.set_mesh` API to enter into a mesh context when using `explicit_axes` API.

What it means

The explicit-axes sharding API needs an active mesh to interpret axis names, but none was found in the current context. JAX stores the current mesh in a context variable set via `jax.set_mesh` (or `mesh_lib.use_abstract_mesh`). Without it, axis names in `axes`/shardings cannot be resolved to device axes.

Source

Thrown at jax/_src/pjit.py:2482

                  in_sharding=None):
  kwargs = dict(axes=axes, in_sharding=in_sharding)
  if f is None:
    return lambda g: _explicit_axes(g, **kwargs)
  return _explicit_axes(f, **kwargs)

def _explicit_axes(fun, *, axes, in_sharding):
  @wraps(fun)
  def decorator(*args, **kwargs):
    if in_sharding is None:
      if "in_sharding" in kwargs:
        _in_sharding = kwargs.pop("in_sharding")
      else:
        raise TypeError("Missing required keyword argument: 'in_sharding'")
    else:
      _in_sharding = in_sharding
    mesh_info = _get_new_mesh(axes, mesh_lib.AxisType.Explicit, 'explicit_axes')
    if mesh_info is None:
      raise ValueError(
          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'
          ' into a mesh context when using `explicit_axes` API.')
    with mesh_lib.use_abstract_mesh(mesh_info.new):
      args = reshard(args, _in_sharding)
      out = fun(*args, **kwargs)
    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(
        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)
    return reshard(out, out_specs)
  return decorator

# -------------------- with_layout_constraint --------------------

def with_layout_constraint(x, layouts):
  x_flat, tree = tree_flatten(x)
  x_avals_flat = [core.shaped_abstractify(x) for x in x_flat]
  layouts_flat = tuple(flatten_axes("with_layout_constraint layouts", tree,
                                    layouts))
  if any(not isinstance(l, Layout) for l in layouts_flat):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap the call in `with jax.set_mesh(mesh):` and move decorated-function invocation inside the block
2. Create the mesh with `jax.make_mesh(...)` before entering the context
3. Ensure the mesh context is still active (not exited) when the decorated function runs

Example fix

// before
f = decorated(fun, in_sharding=P('data'))
f(x)  # no mesh
// after
mesh = jax.make_mesh((8,), ('data',))
with jax.set_mesh(mesh):
  f(x)

import jax
try:
    jax.experimental.mesh.get_abstract_mesh()
    mesh_active = True
except Exception:
    mesh_active = False
if not mesh_active:
    raise RuntimeError('Enter a mesh via jax.set_mesh before calling')

When it happens

Trigger: Using an `explicit_axes` decorator (or `jax.jit` with explicit axes) outside of a `with jax.set_mesh(mesh):` block, or after the mesh context exited (e.g. decorating at module import time but calling later).

Common situations: New explicit/abstract mesh API in recent JAX versions where the old `Mesh(context manager)` pattern no longer applies; decorating functions at module scope while the mesh is only created inside `main()`.

Related errors

• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• callbacks are only supported in spmd computations when all m
• Mesh must be provided for shard_map with checkify.
• Sharding spec {spec} implies that array axis {dim} is partit

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