jax-ml/jax · error · ValueError

Length of sharding.spec ({len(out_s.spec)}) must be equal to

Error message

Length of sharding.spec ({len(out_s.spec)}) must be equal to aval's ndim ({ndim}). Got sharding.spec {out_s.spec}, aval.ndim {ndim} and sharding {out_s}

What it means

jax/_src/core.py canonicalizes a sharding for an abstract value (aval) and requires the length of sharding.spec (a PartitionSpec) to equal the array's ndim. Each element of the spec maps to one array dimension, so a spec with a different length cannot describe how the array is laid out across the mesh.

Source

Thrown at jax/_src/core.py:2295

@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)

  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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make len(PartitionSpec) == array.ndim, adding None for unsharded dims: P('x', None) for a 2-D array
Check arr.ndim and the spec length before attaching the sharding
If reusing shardings across tensors, key them by shape/ndim

Example fix

// before
sh = jax.sharding.NamedSharding(mesh, P('x'))  # 1 entry
x = jax.device_put(x_2d, sh)  # x_2d.ndim == 2

// after
sh = jax.sharding.NamedSharding(mesh, P('x', None))
x = jax.device_put(x_2d, sh)

Defensive patterns

Strategy: validation

Validate before calling

def check_sharding(spec, arr):
    assert len(spec) == arr.ndim, (len(spec), arr.ndim)

Type guard

def spec_matches(spec, ndim): return len(spec) == ndim

Prevention

Always write PartitionSpecs with one entry per dim including None
Validate spec length vs ndim in test fixtures
Key cached shardings by (shape, ndim)

When it happens

Trigger: Passing a NamedSharding whose PartitionSpec has more/fewer entries than the array's rank, e.g. NamedSharding(mesh, P('x')) applied to a 2-D array, or a spec with None placeholders omitted for trailing dims.

Common situations: Reusing a sharding built for one tensor on another of different rank; forgetting that P() entries must cover every dimension including None; mesh/axis refactors changing rank expectations.

Related errors

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