jax-ml/jax · error · ValueError

use_abstract_mesh cannot change the size of the mesh. Got ne

Error message

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}

What it means

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.

Source

Thrown at jax/_src/mesh.py:625

        takes effect if all mesh axes are Explicit. This is temporary until we
        fix the underlying issues.
  """
  __slots__ = ['mesh', 'prev']

  def __init__(self, mesh: AbstractMesh):
    if not isinstance(mesh, AbstractMesh):
      raise ValueError(
          "Expected mesh of type `jax.sharding.AbstractMesh`. Got type:"
          f" {type(mesh)}")
    self.mesh = mesh

  def __enter__(self):
    self.prev = jax_config.abstract_mesh_context_manager.swap_local(self.mesh)
    if (self.prev is not config_ext.unset and
        not self.prev.empty and not self.mesh.empty and
        self.prev.size != self.mesh.size):
      jax_config.abstract_mesh_context_manager.set_local(self.prev)
      raise ValueError(
          "use_abstract_mesh cannot change the size of the mesh. Got new mesh:"
          f" {self.mesh} with size={self.mesh.size} and prev mesh:"
          f" {self.prev} with size={self.prev.size}")

  def __exit__(self, exc_type, exc_value, traceback):
    jax_config.abstract_mesh_context_manager.set_local(self.prev)


def get_abstract_mesh() -> AbstractMesh:
  val = jax_config.abstract_mesh_context_manager.value
  return empty_abstract_mesh if val is None else val

def get_concrete_mesh() -> Mesh:
  val = jax_config.device_context.value
  return empty_concrete_mesh if val is None else val

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the nested abstract mesh's total size equal the outer one (adjust axis sizes or reuse the same size)
2. Exit the outer use_abstract_mesh context before setting a different-size mesh
3. Compute size as prod(axis_sizes) and assert equality before entering

Example fix

# before
with use_abstract_mesh(AbstractMesh(('data','model'), (8, 4))):
  with use_abstract_mesh(AbstractMesh(('data','model'), (16, 4))):  # size 64 != 32
    ...

# after
with use_abstract_mesh(AbstractMesh(('data','model'), (8, 4))):
  with use_abstract_mesh(AbstractMesh(('data','model'), (4, 8))):  # size 32
    ...

import math
new_size = math.prod(new_mesh.axis_sizes)
prev = jax.config.abstract_mesh_context_manager.value
if prev is not None and getattr(prev, 'size', None):
    assert prev.size == new_size, f'{prev.size} != {new_size}'

When it happens

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

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

Related errors

• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• Expected mesh of type `jax.sharding.AbstractMesh`. Got type:
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• callbacks are only supported in spmd computations when all m

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