{"record":{"id":"cbc84bc0ab38dd8b","repo":"jax-ml/jax","slug":"cannot-update-the-mesh-of-the-current-resource-env","errorCode":null,"errorMessage":"Cannot update the mesh of the current resource environment. The new mesh shadows already defined axes {show_axes(overlap)}","messagePattern":"Cannot update the mesh of the current resource environment\\. The new mesh shadows already defined axes (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":55,"sourceCode":"\nzip, unsafe_zip = safe_zip, zip\nconfig_ext = _jax.config\n\nMeshAxisName = Any\nResourceAxisName = Hashable\n\n\ndef show_axes(axes):\n  return \", \".join(sorted(f\"`{a}`\" for a in axes))\n\n\nclass ResourceEnv(NamedTuple):\n  physical_mesh: Mesh\n\n  def with_mesh(self, mesh: Mesh):\n    overlap = set(mesh.axis_names) & (self.resource_axes - set(self.physical_mesh.axis_names))\n    if overlap:\n      raise ValueError(f\"Cannot update the mesh of the current resource \"\n                       f\"environment. The new mesh shadows already defined axes \"\n                       f\"{show_axes(overlap)}\")\n    return self._replace(physical_mesh=mesh)\n\n  @property\n  def physical_resource_axes(self) -> set[ResourceAxisName]:\n    return set(self.physical_mesh.axis_names)\n\n  @property\n  def resource_axes(self) -> set[ResourceAxisName]:\n    return self.physical_resource_axes\n\n  @property\n  def shape(self):\n    return self.physical_mesh.shape\n\n  @property\n  def local_shape(self):","sourceCodeStart":37,"sourceCodeEnd":73,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L37-L73","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"# before\nmesh_inner = jax.sharding.Mesh(devs.reshape(2, 2), ('data', 'model'))\nwith jax.set_mesh(mesh_inner): ...\n# collides if outer context defines 'model'\n\n# after\nmesh_inner = jax.sharding.Mesh(devs.reshape(2, 2), ('data', 'inner_model'))\nwith jax.set_mesh(mesh_inner): ...","handlingStrategy":"validation","validationCode":"outer_axes = {n for n in current_env_axes}  # track axes you've set\nnew_axes = set(mesh.axis_names)\nassert not (new_axes & outer_axes - set(current_physical.axis_names)), 'axis shadowing'","typeGuard":null,"tryCatchPattern":"try:\n    with jax.set_mesh(mesh): train_step()\nexcept ValueError as e:\n    if 'shadows already defined axes' in str(e):\n        mesh = rename_axes(mesh); with jax.set_mesh(mesh): train_step()\n    else: raise","preventionTips":["Use unique axis names per nesting level (prefix them)","Keep a registry of active mesh axes in your framework code"],"tags":["jax","mesh","axis-names","context-manager"],"backgroundTag":"axis-name-collision","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}