jax-ml/jax · error · ValueError

Mesh shape of the input {a.sharding.mesh.shape_tuple} does n

Error message

Mesh shape of the input {a.sharding.mesh.shape_tuple} does not match the mesh shape passed to shard_map  {mesh.shape_tuple} for shape {aval.str_short()}

What it means

When shard_map infers the mesh from its arguments, every input array with a NamedSharding must live on a mesh with the same shape tuple as the mesh passed to (or previously inferred for) shard_map. Different mesh shapes make the sharding ambiguous.

Source

Thrown at jax/_src/shard_map.py:1194

  return mlir.wrap_with_shard_to_full_op(ctx, sx, aval_out, shard_proto,
                                         unspecified)

def _pspec_mhlo_attrs(spec, aval: core.AbstractValue) -> str:
  if isinstance(aval, core.ShapedArray):
    names = _spec_to_names(spec)
    return str(map(names.get, range(aval.ndim)))
  return ''

# Eager evaluation

def get_mesh_from_args(args_flat, mesh):
  for a in args_flat:
    if (hasattr(a, 'sharding') and isinstance(a.sharding, NamedSharding)
        and not a.sharding.mesh.is_scalar):  # pyrefly: ignore[missing-attribute]

      if a.sharding.mesh.shape_tuple != mesh.shape_tuple:
        aval = core.shaped_abstractify(a)
        raise ValueError(
            f"Mesh shape of the input {a.sharding.mesh.shape_tuple} does not"
            " match the mesh shape passed to shard_map "
            f" {mesh.shape_tuple} for shape {aval.str_short()}")
      mesh = a.sharding.mesh
  if isinstance(mesh, AbstractMesh):
    raise ValueError(
        "Please pass `jax.Array`s with a `NamedSharding` as input to"
        " `shard_map` when passing `AbstractMesh` to the mesh argument.")
  assert isinstance(mesh, Mesh)
  return mesh

def _spec_to_vma(spec):
  return frozenset(p for s in spec.partitions if s is not None
                   for p in (s if isinstance(s, tuple) else (s,)))

def _mat_to_spec(mesh, mat):
  return P(order_wrt_mesh(mesh, mat.varying), unreduced=mat.unreduced,
           reduced=mat.reduced, unreduced_kind=mat.unreduced_kind)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Re-create/place input arrays on the same mesh (shape tuple) as the shard_map mesh
Pass the same Mesh object used to build the arrays' NamedShardings to shard_map
Reshard via jax.device_put(arr, NamedSharding(target_mesh, spec)) before the call

Example fix

// before
x = jax.device_put(x, NamedSharding(mesh_a, P('x')))
y = jax.jit(shard_map(f, mesh_b, ...))(x)  # mesh_b shape differs
// after
x = jax.device_put(x, NamedSharding(mesh_b, P('x')))
y = jax.jit(shard_map(f, mesh_b, ...))(x)

Defensive patterns

Strategy: validation

Validate before calling

def same_mesh_shape(args, mesh):
    return all(not hasattr(a, 'sharding') or not isinstance(a.sharding, NamedSharding)
               or a.sharding.mesh.shape_tuple == mesh.shape_tuple for a in args)

Type guard

def args_on_mesh(args, mesh) -> bool:
    return all(getattr(getattr(a, 'sharding', None), 'mesh', mesh).shape_tuple == mesh.shape_tuple for a in args)

Try / catch

try: shard_map(...) except ValueError as e: if 'Mesh shape' in str(e): jax.device_put args onto mesh and retry; else: raise

Prevention

Create one canonical mesh object per process and use it everywhere
Re-put arrays on the current mesh after changing mesh layouts

When it happens

Trigger: Passing arrays placed on a 2x4 mesh to shard_map(mesh=8x1 mesh, ...); typically when arrays come from a different jax.make_mesh/Mesh than the one used in the call.

Common situations: Notebooks or libraries where arrays were created under an earlier/different mesh layout; changing mesh definition between training stages without resharding.

Related errors

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