{"record":{"id":"48a2e3ef92c5eb72","repo":"jax-ml/jax","slug":"mesh-context-manager-is-disabled","errorCode":null,"errorMessage":"Mesh context manager is disabled.","messagePattern":"Mesh context manager is disabled\\.","errorType":"error_code","errorClass":"RuntimeError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":302,"sourceCode":"    axis_types = _normalize_axis_types(axis_names, axis_types, 'Mesh',\n                                       AxisType.Auto)\n    empty = not axis_names and devices_flat[0] is None\n    size = 0 if empty else math.prod(devices.shape)\n    return cls._create(devices_flat, devices.shape, axis_names,\n                       axis_types, size)\n\n  # No __eq__ or __hash__: interned classes use object identity.\n\n  @property\n  def is_scalar(self):\n    return self.size == 1 and not self.axis_names\n\n  def __getnewargs_ex__(self):\n    return (self.devices, self.axis_names, self.axis_types), {}\n\n  def __enter__(self):\n    if jax_config.disallow_mesh_context_manager.value:\n      raise RuntimeError(\"Mesh context manager is disabled.\")\n    warnings.warn(\n        \"`with mesh:` context manager has been deprecated. Please use `with\"\n        \" jax.set_mesh(mesh):` instead.\",\n        category=DeprecationWarning, stacklevel=2)\n    new_env = thread_resources.stack[-1].with_mesh(self)\n    thread_resources.stack.append(new_env)\n    thread_resources.env = new_env\n    jax_config.mesh_context_manager.set_local(\n        tuple(t.physical_mesh for t in thread_resources.stack\n              if not t.physical_mesh.empty))\n    return self\n\n  def __exit__(self, exc_type, exc_value, traceback):\n    thread_resources.stack.pop()\n    thread_resources.env = thread_resources.stack[-1]\n    jax_config.mesh_context_manager.set_local(\n        tuple(t.physical_mesh for t in thread_resources.stack\n              if not t.physical_mesh.empty))","sourceCodeStart":284,"sourceCodeEnd":320,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L284-L320","documentation":"Mesh.__enter__ raises RuntimeError when the config flag jax_disallow_mesh_context_manager (typically via JAX_DISALLOW_MESH_CONTEXT_MANAGER or jax.config) is set. It also emits a DeprecationWarning: the `with mesh:` pattern is deprecated in favor of jax.set_mesh.","triggerScenarios":"Using `with mesh:` while the disallow flag is on — e.g. inside libraries that set the flag to enforce the new API, or after jax.config.update('jax_disallow_mesh_context_manager', True).","commonSituations":"JAX deprecating the context-manager mesh; frameworks (e.g. newer Flax/JAX releases or test suites) enabling the flag globally; user code relying on implicit mesh context from `with mesh:`.","solutions":["Replace `with mesh:` with `with jax.set_mesh(mesh):`","Or use explicit sharding arguments (NamedSharding) instead of ambient mesh context","Unset JAX_DISALLOW_MESH_CONTEXT_MANAGER if you must keep old code temporarily"],"exampleFix":"# before\nwith mesh:\n  out = jax.jit(f)(x)\n\n# after\nwith jax.set_mesh(mesh):\n  out = jax.jit(f)(x)","handlingStrategy":"validation","validationCode":"import jax\nassert not jax.config.jax_disallow_mesh_context_manager.value, 'with mesh: disabled'","typeGuard":null,"tryCatchPattern":"try:\n    with mesh: run()\nexcept RuntimeError as e:\n    if 'disabled' in str(e):\n        with jax.set_mesh(mesh): run()\n    else: raise","preventionTips":["Migrate all `with mesh:` to `with jax.set_mesh(mesh):`","Prefer explicit sharding args over ambient mesh context"],"tags":["jax","mesh","deprecation","context-manager","config-flag"],"backgroundTag":"deprecated-api-disabled","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}