jax-ml/jax · error · ValueError

Sharding spec {spec} implies that array axis {dim} is partit

Error message

Sharding spec {spec} implies that array axis {dim} is partitioned {size} times, but does not evenly divide the dimension size {sh}. Got shape: {shape} and sharding {sharding}

What it means

Raised when building a sharded array type: an axis of the array is partitioned across a mesh dimension product that does not evenly divide that axis's size, producing an impossible/remainder sharding.

Source

Thrown at jax/_src/core.py:2273

  else:
    out = sharding.update(spec=modify_spec_for_auto_manual(
        sharding.spec, sharding.mesh))
  if config.remove_size_one_mesh_axis_from_type.value:
    out = out.update(spec=ns.remove_size_one_mesh_axis_from_spec(out.spec, out.mesh))
  if len(out.spec) != ndim:
    out = _make_lengths_same(out, ndim)
  return out

def _check_divisibility(sharding, shape):
  mesh = sharding.mesh
  for dim, (spec, sh) in enumerate(zip(sharding.spec.partitions, shape)):
    if spec is None:
      continue
    spec = spec if isinstance(spec, tuple) else (spec,)
    size = math.prod(mesh.shape[s] for s in spec)
    _, remainder = divmod(sh, size)
    if remainder != 0:
      raise ValueError(
          f"Sharding spec {spec} implies that array axis {dim} is partitioned"
          f" {size} times, but does not evenly divide the dimension size {sh}."
          f" Got shape: {shape} and sharding {sharding}")

@cache(max_size=4096,
       trace_context_in_key=lambda: config.remove_size_one_mesh_axis_from_type.value)
def get_sharding(sharding, shape):
  """Modifies and checks the sharding.

  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)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Adjust the PartitionSpec so the product of mesh dims for each axis divides that axis's size (or use None to replicate)
2. Pad or reshape the array so the sharded dimension is divisible by the mesh sub-shape
3. Resize/reorder the mesh (mesh_shape) so each partitioned axis divides evenly
4. Compute divisibility programmatically before constructing the sharding (see validation code)

Example fix

# before
mesh = Mesh(jax.devices(), ('x',))
sharding = NamedSharding(mesh, P('x',))  # shape (3, 4) over 8 devices -> error

# after
mesh = Mesh(jax.devices()[:4], ('x',))  # axis 0 size 4, wait 3 not divisible either
# correct: shard axis of size 4 over 4 devices
sharding = NamedSharding(mesh, P(None, 'x'))  # shape (3, 4), shard axis 1 of size 4 over 4

import math
def sharding_divides(shape, mesh_shape, spec):
    for dim, s in enumerate(spec):
        if s is None: continue
        s = (s,) if not isinstance(s, tuple) else s
        size = math.prod(mesh_shape[axis] for axis in s)
        if shape[dim] % size != 0:
            return False
    return True
assert sharding_divides(x.shape, mesh.shape, P('data',).specs)

When it happens

Trigger: Using NamedSharding/PartitionSpec (or GSPMD sharding annotations) where the product of mesh dimension sizes assigned to an array axis does not divide that axis's length, e.g., sharding a size-3 axis over a 2-device mesh dimension.

Common situations: Mismatch between array shape and jax.sharding.Mesh shape; leftover partitions from an older mesh; using pmap with more devices than a dimension's size; sharding tiny axes (size 1 or primes) across multi-device mesh axes.

Related errors

• Reduced axes can only refer to mesh axes that is of type `Ex
• 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.

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