{"record":{"id":"49e12848c7edbde3","repo":"jax-ml/jax","slug":"context-mesh-cannot-be-empty-please-use-jax-set","errorCode":null,"errorMessage":"Context mesh cannot be empty. Please use `jax.set_mesh` API to enter into a mesh context when using `explicit_axes` API.","messagePattern":"Context mesh cannot be empty\\. Please use `jax\\.set_mesh` API to enter into a mesh context when using `explicit_axes` API\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pjit.py","lineNumber":2482,"sourceCode":"                  in_sharding=None):\n  kwargs = dict(axes=axes, in_sharding=in_sharding)\n  if f is None:\n    return lambda g: _explicit_axes(g, **kwargs)\n  return _explicit_axes(f, **kwargs)\n\ndef _explicit_axes(fun, *, axes, in_sharding):\n  @wraps(fun)\n  def decorator(*args, **kwargs):\n    if in_sharding is None:\n      if \"in_sharding\" in kwargs:\n        _in_sharding = kwargs.pop(\"in_sharding\")\n      else:\n        raise TypeError(\"Missing required keyword argument: 'in_sharding'\")\n    else:\n      _in_sharding = in_sharding\n    mesh_info = _get_new_mesh(axes, mesh_lib.AxisType.Explicit, 'explicit_axes')\n    if mesh_info is None:\n      raise ValueError(\n          'Context mesh cannot be empty. Please use `jax.set_mesh` API to enter'\n          ' into a mesh context when using `explicit_axes` API.')\n    with mesh_lib.use_abstract_mesh(mesh_info.new):\n      args = reshard(args, _in_sharding)\n      out = fun(*args, **kwargs)\n    out_specs = tree_map(lambda o: core.modify_spec_for_auto_manual(\n        core.typeof(o).sharding.spec, mesh_lib.get_abstract_mesh()), out)\n    return reshard(out, out_specs)\n  return decorator\n\n# -------------------- with_layout_constraint --------------------\n\ndef with_layout_constraint(x, layouts):\n  x_flat, tree = tree_flatten(x)\n  x_avals_flat = [core.shaped_abstractify(x) for x in x_flat]\n  layouts_flat = tuple(flatten_axes(\"with_layout_constraint layouts\", tree,\n                                    layouts))\n  if any(not isinstance(l, Layout) for l in layouts_flat):","sourceCodeStart":2464,"sourceCodeEnd":2500,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pjit.py#L2464-L2500","documentation":"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.","triggerScenarios":"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).","commonSituations":"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()`.","solutions":["Wrap the call in `with jax.set_mesh(mesh):` and move decorated-function invocation inside the block","Create the mesh with `jax.make_mesh(...)` before entering the context","Ensure the mesh context is still active (not exited) when the decorated function runs"],"exampleFix":"// before\nf = decorated(fun, in_sharding=P('data'))\nf(x)  # no mesh\n// after\nmesh = jax.make_mesh((8,), ('data',))\nwith jax.set_mesh(mesh):\n  f(x)","handlingStrategy":"validation","validationCode":"import jax\ntry:\n    jax.experimental.mesh.get_abstract_mesh()\n    mesh_active = True\nexcept Exception:\n    mesh_active = False\nif not mesh_active:\n    raise RuntimeError('Enter a mesh via jax.set_mesh before calling')","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Always enter mesh via `with jax.set_mesh(mesh):` around decorated calls","Create the mesh once at startup with jax.make_mesh and pass context explicitly"],"tags":["jax","mesh","sharding","context"],"backgroundTag":"missing-context-manager","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}