jax-ml/jax · error · RuntimeError

Mesh context manager is disabled.

Error message

Mesh context manager is disabled.

What it means

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.

Source

Thrown at jax/_src/mesh.py:302

    axis_types = _normalize_axis_types(axis_names, axis_types, 'Mesh',
                                       AxisType.Auto)
    empty = not axis_names and devices_flat[0] is None
    size = 0 if empty else math.prod(devices.shape)
    return cls._create(devices_flat, devices.shape, axis_names,
                       axis_types, size)

  # No __eq__ or __hash__: interned classes use object identity.

  @property
  def is_scalar(self):
    return self.size == 1 and not self.axis_names

  def __getnewargs_ex__(self):
    return (self.devices, self.axis_names, self.axis_types), {}

  def __enter__(self):
    if jax_config.disallow_mesh_context_manager.value:
      raise RuntimeError("Mesh context manager is disabled.")
    warnings.warn(
        "`with mesh:` context manager has been deprecated. Please use `with"
        " jax.set_mesh(mesh):` instead.",
        category=DeprecationWarning, stacklevel=2)
    new_env = thread_resources.stack[-1].with_mesh(self)
    thread_resources.stack.append(new_env)
    thread_resources.env = new_env
    jax_config.mesh_context_manager.set_local(
        tuple(t.physical_mesh for t in thread_resources.stack
              if not t.physical_mesh.empty))
    return self

  def __exit__(self, exc_type, exc_value, traceback):
    thread_resources.stack.pop()
    thread_resources.env = thread_resources.stack[-1]
    jax_config.mesh_context_manager.set_local(
        tuple(t.physical_mesh for t in thread_resources.stack
              if not t.physical_mesh.empty))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Replace `with mesh:` with `with jax.set_mesh(mesh):`
2. Or use explicit sharding arguments (NamedSharding) instead of ambient mesh context
3. Unset JAX_DISALLOW_MESH_CONTEXT_MANAGER if you must keep old code temporarily

Example fix

# before
with mesh:
  out = jax.jit(f)(x)

# after
with jax.set_mesh(mesh):
  out = jax.jit(f)(x)

import jax
assert not jax.config.jax_disallow_mesh_context_manager.value, 'with mesh: disabled'

try:
    with mesh: run()
except RuntimeError as e:
    if 'disabled' in str(e):
        with jax.set_mesh(mesh): run()
    else: raise

When it happens

Trigger: 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).

Common situations: 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:`.

Related errors

• Cannot update the mesh of the current resource environment.
• When `check_vma=True` on `jax.shard_map`, `manual_axis_type`
• The return value of the policies should be a boolean. Got: {
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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