jax-ml/jax · error · ValueError

AbstractMesh should be of the same size across all shardings

Error message

AbstractMesh should be of the same size across all shardings. Got {abstract_mesh} and {sh.mesh}

What it means

When a computation mixes NamedShardings over an AbstractMesh, JAX requires all such meshes to have the same total size (number of devices). If shardings in the same lowering carry AbstractMeshes of differing sizes, device assignment is ambiguous and JAX raises during _get_and_check_device_assignment (used by lower_sharding_computation).

Source

Thrown at jax/_src/interpreters/pxla.py:567

def _get_and_check_device_assignment(
    shardings: Iterable[ShardingInfo],
    ctx_mesh: Mesh | AbstractMesh,
) -> tuple[xc.Client, tuple[xc.Device, ...] | None, int]:
  first_sharding_info = None
  abstract_mesh = (
      ctx_mesh if not ctx_mesh.empty and isinstance(ctx_mesh, AbstractMesh)
      else None)
  any_concrete_sharding = (
      True if not ctx_mesh.empty and isinstance(ctx_mesh, Mesh) else False)

  for sh, s_type, source_info in shardings:
    if isinstance(sh, UnspecifiedValue):
      continue
    elif isinstance(sh, NamedSharding) and isinstance(sh.mesh, AbstractMesh):
      if (abstract_mesh is not None and not sh.mesh.empty and
          abstract_mesh.size != sh.mesh.size):
        raise ValueError("AbstractMesh should be of the same size across all "
                         f"shardings. Got {abstract_mesh} and {sh.mesh}")
      abstract_mesh = sh.mesh
    else:
      any_concrete_sharding = True
      arr_device_assignment = sh._device_assignment
      if first_sharding_info is None:
        first_sharding_info = (arr_device_assignment, s_type, source_info)
      if ctx_mesh.empty:
        if first_sharding_info[0] != arr_device_assignment:
          raise stages.DeviceAssignmentMismatchError([
              stages.DeviceAssignmentMismatch(*first_sharding_info),
              stages.DeviceAssignmentMismatch(
                  arr_device_assignment, s_type, source_info)])
      elif isinstance(ctx_mesh, AbstractMesh):
        if ctx_mesh.size != len(arr_device_assignment):
          raise stages.DeviceAssignmentMismatchError([
              stages.DeviceAssignmentMismatch(
                  ctx_mesh.size, stages.MismatchType.CONTEXT_DEVICES, None),

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make all AbstractMeshes used for in/out shardings the same size (same total device count)
2. Derive all shardings from a single shared Mesh/AbstractMesh object
3. Check mesh.empty cases: ensure you aren't mixing an empty AbstractMesh placeholder with a sized one inconsistently

Example fix

# before
mesh_a = Mesh(devices8, axis_names=('dp',))
mesh_b = Mesh(devices4, axis_names=('dp',))
jit(f, in_shardings=NamedSharding(mesh_a, P('dp')), out_shardings=NamedSharding(mesh_b, P('dp')))

# after
mesh = Mesh(devices8, axis_names=('dp',))
jit(f, in_shardings=NamedSharding(mesh, P('dp')), out_shardings=NamedSharding(mesh, P('dp')))

meshes = [s.mesh for s in shardings if isinstance(s, NamedSharding) and isinstance(s.mesh, AbstractMesh)]
sizes = {m.size for m in meshes if not m.empty}
assert len(sizes) <= 1, f'AbstractMesh sizes differ: {sizes}'

When it happens

Trigger: Constructing a jitted function (pjit) where different arguments' NamedShardings use AbstractMesh objects with different total sizes — e.g. one mesh 2x4 and another 8x1, or one empty and one populated inconsistently; mixing out_shardings and in_shardings with different abstract meshes.

Common situations: Building shardings from multiple user-created Mesh/AbstractMesh objects; refactoring mesh shapes mid-experiment; passing an out_sharding from an old mesh with new in_shardings.

Related errors

• AbstractMesh size: {abstract_mesh.size} does not match the d
• Mesh of an aval must be an AbstractMesh. Got {out_s.mesh} of
• Passing lowering_platforms via jax.export or jit(f).trace(*a
• Missing required keyword argument: 'in_sharding'
• {self} implies that array axis {dim} is partitioned {p} time

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