{"record":{"id":"2d7dd8b8c62b9e81","repo":"jax-ml/jax","slug":"the-context-mesh-cannot-be-empty-use-jax-set-mes","errorCode":null,"errorMessage":"The context mesh cannot be empty. Use `jax.set_mesh(mesh)` to enter into a mesh context","messagePattern":"The context mesh cannot be empty\\. Use `jax\\.set_mesh\\(mesh\\)` to enter into a mesh context","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/shard_map.py","lineNumber":368,"sourceCode":"    except _RepError as e:\n      fails, out_tree, = e.args\n      msg = _inout_vma_error(f, mesh, out_tree, out_specs, fails)\n      raise ValueError(msg) from None\n    return out_ft.unflatten()\n  return cast(F, wrapped)\n\n\ndef _axes_to_pspec(axis_name, axis):\n  if axis is None:\n    return P()\n  return P(*[None] * axis + [axis_name])\n\n\ndef _shmap_checks(mesh, axis_names, in_specs, out_specs, _smap):\n  if mesh is None:\n    mesh = get_abstract_mesh()\n    if mesh.empty:\n      raise ValueError(\n          \"The context mesh cannot be empty. Use\"\n          \" `jax.set_mesh(mesh)` to enter into a mesh context\")\n  else:\n    ctx_mesh = get_abstract_mesh()\n    if not ctx_mesh.empty and mesh.abstract_mesh != ctx_mesh:\n      raise ValueError(\n          f\"The context mesh {ctx_mesh} should match the mesh passed to\"\n          f\" shard_map {mesh}\")\n\n  if not isinstance(mesh, (Mesh, AbstractMesh)):\n    raise TypeError(\"shard_map requires a `jax.sharding.Mesh` or a \"\n                    \"`jax.sharding.AbstractMesh` instance for its \"\n                    f\"second argument, but got {mesh} of type {type(mesh)}.\")\n  if mesh.empty:\n    raise ValueError(f\"shard_map requires a non-empty mesh. Got {mesh}\")\n\n  mesh_axis_names_wo_vmap = (\n      frozenset(mesh.axis_names) - core.get_axis_env().explicit_mesh_axis_names","sourceCodeStart":350,"sourceCodeEnd":386,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/shard_map.py#L350-L386","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","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)"],"exampleFix":"// before\nout = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))(x)  # no mesh, no context\n\n// after\nout = jax.shard_map(f, mesh=mesh, in_specs=P('i'), out_specs=P('i'))(x)\n# or\nwith jax.set_mesh(mesh):\n    out = jax.shard_map(f, in_specs=P('i'), out_specs=P('i'))(x)","handlingStrategy":"validation","validationCode":"from jax.sharding import AbstractMesh\n\ndef get_active_mesh(explicit=None):\n    m = explicit\n    if m is None:\n        m = jax.sharding.get_abstract_mesh()\n        if m.empty:\n            raise ValueError('no context mesh; pass mesh= or use jax.set_mesh')\n    return m","typeGuard":null,"tryCatchPattern":null,"preventionTips":["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"],"tags":["jax","shard-map","mesh","context","distributed"],"backgroundTag":"missing-mesh-context","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}