jax-ml/jax · error · ValueError

Cannot update the mesh of the current resource environment.

Error message

Cannot update the mesh of the current resource environment. The new mesh shadows already defined axes {show_axes(overlap)}

What it means

ResourceEnv.with_mesh computes the overlap between the new mesh's axis names and resource axes defined outside the current physical mesh (e.g. from an outer namedsharding or abstract mesh context). Overlapping names would silently shadow those axes, so it raises ValueError. This backs the deprecated `with mesh:` and jax.set_mesh context managers.

Source

Thrown at jax/_src/mesh.py:55

zip, unsafe_zip = safe_zip, zip
config_ext = _jax.config

MeshAxisName = Any
ResourceAxisName = Hashable


def show_axes(axes):
  return ", ".join(sorted(f"`{a}`" for a in axes))


class ResourceEnv(NamedTuple):
  physical_mesh: Mesh

  def with_mesh(self, mesh: Mesh):
    overlap = set(mesh.axis_names) & (self.resource_axes - set(self.physical_mesh.axis_names))
    if overlap:
      raise ValueError(f"Cannot update the mesh of the current resource "
                       f"environment. The new mesh shadows already defined axes "
                       f"{show_axes(overlap)}")
    return self._replace(physical_mesh=mesh)

  @property
  def physical_resource_axes(self) -> set[ResourceAxisName]:
    return set(self.physical_mesh.axis_names)

  @property
  def resource_axes(self) -> set[ResourceAxisName]:
    return self.physical_resource_axes

  @property
  def shape(self):
    return self.physical_mesh.shape

  @property
  def local_shape(self):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Rename the inner mesh's axis names so they don't collide with outer resource axes
Exit the outer mesh/abstract-mesh context before entering the new mesh
Audit all enclosing set_mesh / use_abstract_mesh blocks for axis names already defined

Example fix

# before
mesh_inner = jax.sharding.Mesh(devs.reshape(2, 2), ('data', 'model'))
with jax.set_mesh(mesh_inner): ...
# collides if outer context defines 'model'

# after
mesh_inner = jax.sharding.Mesh(devs.reshape(2, 2), ('data', 'inner_model'))
with jax.set_mesh(mesh_inner): ...

Defensive patterns

Strategy: validation

Validate before calling

outer_axes = {n for n in current_env_axes}  # track axes you've set
new_axes = set(mesh.axis_names)
assert not (new_axes & outer_axes - set(current_physical.axis_names)), 'axis shadowing'

Try / catch

try:
    with jax.set_mesh(mesh): train_step()
except ValueError as e:
    if 'shadows already defined axes' in str(e):
        mesh = rename_axes(mesh); with jax.set_mesh(mesh): train_step()
    else: raise

Prevention

Use unique axis names per nesting level (prefix them)
Keep a registry of active mesh axes in your framework code

When it happens

Trigger: Entering a mesh context manager while an outer context already defines an axis name that is not on the current physical mesh — e.g. nesting set_mesh calls where the inner mesh reuses an axis name ('replica', 'data', 'model') defined by an outer abstract mesh.

Common situations: Nesting mesh contexts in multi-level pipelines (e.g. physical mesh inside an abstract mesh context); refactoring code that reuses axis names like 'batch' or 'model' at multiple nesting levels; migration to jax.set_mesh exposing previously hidden conflicts.

Related errors

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