jax-ml/jax · error · ValueError

Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of

Error message

Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of type {type(out_s.mesh)}

What it means

When an abstract value carries a sharding, JAX requires its mesh to be an AbstractMesh (the named-axis mesh used inside traced/computation avals). A concrete runtime mesh or other object type is rejected because aval-level shardings only describe logical named-axis layouts.

Source

Thrown at jax/_src/core.py:2300

  Some modifications/checks include:
    * Making the length of specs the same as ndim
    * If a mesh axis is mentioned in pspec is Auto/Manual, replace it with None
    * Checking for len(spec)-ndim match
    * Checking if the mesh is an AbstractMesh.
  """
  ndim = len(shape)
  if sharding is None:
    return _empty_sharding(ndim)

  out_s = _maybe_modify_sharding(sharding, ndim)
  if len(out_s.spec) != ndim:
    raise ValueError(
        f"Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's"
        f" ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and"
        f" sharding {out_s}")
  if not isinstance(out_s.mesh, mesh_lib.AbstractMesh):
    raise ValueError("Mesh of an aval must be an AbstractMesh. "
                     f"Got {out_s.mesh} of type {type(out_s.mesh)}")
  _check_divisibility(out_s, shape)
  if out_s.memory_kind is not None:
    raise ValueError(
        "sharding with memory_kind is not allowed. Please use `jax.device_put`"
        f" to transfer to different memory spaces. Got {sharding=}")
  return out_s


@cache(max_size=4096,
       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)
def get_mat(mat, mesh):
  if mesh.empty:
    assert mat.empty, mat
    return mat

  axis_env = get_axis_env()
  in_axis_env = lambda i: axis_env.axis_exists(i) and i not in mesh._name_to_type

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.make_mesh(..., abstract=True) or an AbstractMesh when building shardings destined for avals
2. Don't attach runtime NamedSharding objects to abstract values; only user-level APIs (device_put, jit in_sharding) should
3. Upgrade JAX so mesh-handling matches the current API surface

Example fix

// before
mesh = jax.make_mesh((4,), ('x',))  # concrete mesh
aval_sharding = NamedSharding(mesh, P('x'))  # then used on an aval

// after
amesh = jax.make_mesh((4,), ('x',), abstract=True)
aval_sharding = NamedSharding(amesh, P('x'))

from jax._src import mesh_lib

def is_abstract_mesh_sharding(s):
    return isinstance(getattr(s, 'mesh', None), mesh_lib.AbstractMesh)

try:
    aval_with_sharding(s)
except ValueError as e:
    if 'AbstractMesh' in str(e):
        s = rebuild_with_abstract_mesh(s)

When it happens

Trigger: Constructing avals or calling internal APIs like get_sharding/aval constructors with a sharding whose .mesh is a concrete Mesh or a user object; passing a NamedSharding created from a physical Mesh into aval-level code paths.

Common situations: Mixing jax.sharding.Mesh (concrete) with jax.mesh_utils / AbstractMesh APIs; version upgrades where internal invariants on aval shardings became enforced; writing custom primitives that attach shardings to avals.

Related errors

• 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
• Mesh must be provided for shard_map with checkify.
• Sharding spec {spec} implies that array axis {dim} is partit

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