jax-ml/jax · error · ValueError

Expected mesh of type `jax.sharding.AbstractMesh`. Got type:

Error message

Expected mesh of type `jax.sharding.AbstractMesh`. Got type: {type(mesh)}

What it means

The use_abstract_mesh context manager validates its argument with isinstance(mesh, AbstractMesh) and raises ValueError otherwise. Note the target type is AbstractMesh specifically — passing a concrete Mesh is also an error here (in most versions).

Source

Thrown at jax/_src/mesh.py:614

    @jax.jit
    def f(x):
      return x * 2

    with jax.sharding.use_abstract_mesh(abstract_mesh):
      # Note: `f` will be traced and lowered for TPU platform.
      f.trace(inp).lower()
      # Note: `f` will be traced for TPU and lowered for CPU.
      f.trace(inp).lower(lowering_platforms=('cpu',))

  Note: In the example above, setting the abstract mesh at the top level only
        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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Construct an AbstractMesh(axis_names, axis_sizes) and pass that
2. If you meant to set a concrete mesh, use jax.set_mesh / Mesh context instead
3. Check isinstance before entering the context in generic code

Example fix

# before
with use_abstract_mesh(jax.sharding.Mesh(devs, ('data','model'))):
  ...

# after
from jax.sharding import AbstractMesh
with use_abstract_mesh(AbstractMesh(('data','model'), (8, 4))):
  ...

from jax.sharding import AbstractMesh
assert isinstance(mesh, AbstractMesh), 'use_abstract_mesh needs AbstractMesh'

from jax.sharding import AbstractMesh
def is_abstract_mesh(m) -> bool:
    return isinstance(m, AbstractMesh)

When it happens

Trigger: Calling jax.experimental.../mesh use_abstract_mesh(Mesh(...)) with a concrete Mesh, None, or a duck-typed object.

Common situations: Confusion over the dual mesh types introduced for autosharding: developers pass a physical Mesh where an abstract (name/axis-only) mesh is expected; migrating configs where the same variable holds either type.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• axis_types passed to {name} must be of type `jax.sharding.Ax
• use_abstract_mesh cannot change the size of the mesh. Got ne
• 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/3f6ed0e306537f02. Report an issue: GitHub.