jax-ml/jax · error · ValueError

AbstractMesh size: {abstract_mesh.size} does not match the d

Error message

AbstractMesh size: {abstract_mesh.size} does not match the device assignment size: {len(device_assignment)}

What it means

When a computation mixes concrete shardings with NamedShardings over an AbstractMesh, JAX takes the device assignment from the concrete sharding and checks it against the AbstractMesh's total size; they must match. A mismatch means the sharded computation would be built for a different device set than the mesh declares.

Source

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

                  ctx_mesh._flat_devices_tuple,
                  stages.MismatchType.CONTEXT_DEVICES, None),
              stages.DeviceAssignmentMismatch(
                  arr_device_assignment, s_type, source_info)])

  device_assignment: tuple[xc.Device, ...]
  if (first_sharding_info is None and not ctx_mesh.empty and
      isinstance(ctx_mesh, Mesh)):
    device_assignment = ctx_mesh._flat_devices_tuple
  elif first_sharding_info is None:
    device_assignment = (get_default_device(),)
  else:
    device_assignment = first_sharding_info[0]  # pyrefly: ignore[bad-assignment]

  backend = xb.get_device_backend(device_assignment[0])

  if (any_concrete_sharding and abstract_mesh is not None and
      len(device_assignment) != abstract_mesh.size):
    raise ValueError(
        f"AbstractMesh size: {abstract_mesh.size} does not match the"
        f" device assignment size: {len(device_assignment)}")

  if any_concrete_sharding or abstract_mesh is None:
    return backend, device_assignment, len(device_assignment)
  else:
    return backend, None, abstract_mesh.size

MaybeSharding = JSharding | UnspecifiedValue


def prune_unused_inputs(
    jaxpr: core.Jaxpr,
) -> tuple[core.Jaxpr, set[int], set[int]]:
  used_outputs = [True] * len(jaxpr.outvars)
  new_jaxpr, used_consts, used_inputs = pe.dce_jaxpr_consts(jaxpr, used_outputs)
  kept_const_idx = {i for i, b in enumerate(used_consts) if b}
  kept_var_idx = {i for i, b in enumerate(used_inputs) if b}

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make the concrete device_assignment length equal the AbstractMesh size (same device set)
2. Prefer all-abstract or all-concrete shardings rather than mixing
3. If mixing intentionally, construct the AbstractMesh over exactly the devices in the concrete assignment

Example fix

# before
mesh = AbstractMesh((8,), axis_names=('dp',))  # size 8
devs = jax.devices()[:4]
jit(f, in_shardings=(GSPMDSharding(devs, ...), NamedSharding(mesh, P('dp'))))

# after
mesh = AbstractMesh((4,), axis_names=('dp',))  # matches len(devs)
jit(f, in_shardings=(GSPMDSharding(devs, ...), NamedSharding(mesh, P('dp'))))

abstract_sizes = [s.mesh.size for s in shardings if isinstance(s, NamedSharding) and isinstance(s.mesh, AbstractMesh)]
concrete = [s._device_assignment for s in shardings if not isinstance(s, NamedSharding)]
if abstract_sizes and concrete:
    assert all(len(d) == abstract_sizes[0] for d in concrete)

When it happens

Trigger: pjit with some args having concrete device shardings (e.g. SingleDeviceSharding, GSPMDSharding with a device list of N devices) and others using NamedSharding over an AbstractMesh of size M != N.

Common situations: Mixing concrete device assignments (e.g. from jax.device_put results or explicit GSPMDSharding) with abstract-mesh-based NamedShardings on other arguments/outputs; partial migration to AbstractMesh-based APIs.

Related errors

• AbstractMesh should be of the same size across all shardings
• 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/4b94710e99745cc3. Report an issue: GitHub.