{"record":{"id":"9396bde986219226","repo":"jax-ml/jax","slug":"use-abstract-mesh-cannot-change-the-size-of-the-me","errorCode":null,"errorMessage":"use_abstract_mesh cannot change the size of the mesh. Got new mesh: {self.mesh} with size={self.mesh.size} and prev mesh: {self.prev} with size={self.prev.size}","messagePattern":"use_abstract_mesh cannot change the size of the mesh\\. Got new mesh: (.+?) with size=(.+?) and prev mesh: (.+?) with size=(.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/mesh.py","lineNumber":625,"sourceCode":"        takes effect if all mesh axes are Explicit. This is temporary until we\n        fix the underlying issues.\n  \"\"\"\n  __slots__ = ['mesh', 'prev']\n\n  def __init__(self, mesh: AbstractMesh):\n    if not isinstance(mesh, AbstractMesh):\n      raise ValueError(\n          \"Expected mesh of type `jax.sharding.AbstractMesh`. Got type:\"\n          f\" {type(mesh)}\")\n    self.mesh = mesh\n\n  def __enter__(self):\n    self.prev = jax_config.abstract_mesh_context_manager.swap_local(self.mesh)\n    if (self.prev is not config_ext.unset and\n        not self.prev.empty and not self.mesh.empty and\n        self.prev.size != self.mesh.size):\n      jax_config.abstract_mesh_context_manager.set_local(self.prev)\n      raise ValueError(\n          \"use_abstract_mesh cannot change the size of the mesh. Got new mesh:\"\n          f\" {self.mesh} with size={self.mesh.size} and prev mesh:\"\n          f\" {self.prev} with size={self.prev.size}\")\n\n  def __exit__(self, exc_type, exc_value, traceback):\n    jax_config.abstract_mesh_context_manager.set_local(self.prev)\n\n\ndef get_abstract_mesh() -> AbstractMesh:\n  val = jax_config.abstract_mesh_context_manager.value\n  return empty_abstract_mesh if val is None else val\n\ndef get_concrete_mesh() -> Mesh:\n  val = jax_config.device_context.value\n  return empty_concrete_mesh if val is None else val\n","sourceCodeStart":607,"sourceCodeEnd":641,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/mesh.py#L607-L641","documentation":"use_abstract_mesh.__enter__ swaps the thread-local abstract mesh but forbids changing mesh size while an abstract mesh is already set: if both previous and new meshes are non-empty and their sizes differ, it restores the previous mesh and raises ValueError. Consistent size is required because sharding computations assume a fixed world size per context chain.","triggerScenarios":"Nesting use_abstract_mesh contexts where the inner AbstractMesh has a different total size (product of axis sizes) than the outer one — e.g. outer (8,4) size 32, inner (16,4) size 64.","commonSituations":"Multi-stage LLM training pipelines (e.g. MaxText/JAX-Toolbox style) entering an abstract mesh per stage; updating a mesh config (FSDP or TP degree) in only one nested context; refactoring that changes one axis size and silently changes total size.","solutions":["Make the nested abstract mesh's total size equal the outer one (adjust axis sizes or reuse the same size)","Exit the outer use_abstract_mesh context before setting a different-size mesh","Compute size as prod(axis_sizes) and assert equality before entering"],"exampleFix":"# before\nwith use_abstract_mesh(AbstractMesh(('data','model'), (8, 4))):\n  with use_abstract_mesh(AbstractMesh(('data','model'), (16, 4))):  # size 64 != 32\n    ...\n\n# after\nwith use_abstract_mesh(AbstractMesh(('data','model'), (8, 4))):\n  with use_abstract_mesh(AbstractMesh(('data','model'), (4, 8))):  # size 32\n    ...","handlingStrategy":"validation","validationCode":"import math\nnew_size = math.prod(new_mesh.axis_sizes)\nprev = jax.config.abstract_mesh_context_manager.value\nif prev is not None and getattr(prev, 'size', None):\n    assert prev.size == new_size, f'{prev.size} != {new_size}'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Assert prod(axis_sizes) equality before nesting use_abstract_mesh","Define one world size per context chain in your framework"],"tags":["jax","mesh","abstract-mesh","size-mismatch"],"backgroundTag":"world-size-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}