jax-ml/jax · error · ValueError

partitions cannot overlap with reduced axes passed to Partit

Error message

partitions cannot overlap with reduced axes passed to PartitionSpec. Got partitions: {partitions} and reduced axes: {reduced}

What it means

Axes in the positional partitions of a PartitionSpec cannot also appear in the reduced set; a reduced axis is fully summed and cannot simultaneously define how a dimension is sharded.

Source

Thrown at jax/_src/partition_spec.py:55

  if unreduced_kind is not None and not isinstance(unreduced_kind, UnreducedKind):
      raise TypeError(
          "Expected unreduced_kind to be of type `jax.sharding.UnreducedKind`"
          f" but got {type(unreduced_kind)}")
  if not unreduced and unreduced_kind is not None:
    raise ValueError(
        "`unreduced_kind` should be `None` when `unreduced` is an empty set."
        f" Got {unreduced_kind=} and {unreduced=}")

  for partition in partitions:
    partition = partition if isinstance(partition, tuple) else (partition,)
    for p in partition:
      if p in unreduced:
        raise ValueError(
            "partitions cannot overlap with unreduced axes passed to"
            f" PartitionSpec. Got partitions: {partitions} and unreduced axes:"
            f" {unreduced}")
      if p in reduced:
        raise ValueError(
            "partitions cannot overlap with reduced axes passed to"
            f" PartitionSpec. Got partitions: {partitions} and reduced axes:"
            f" {reduced}")

def _get_ur_str(unreduced, reduced):
  if unreduced and reduced:
    return f"unreduced={set(unreduced)!r}, reduced={set(reduced)!r}"
  elif unreduced and not reduced:
    return f"unreduced={set(unreduced)!r}"
  elif not unreduced and reduced:
    return f"reduced={set(reduced)!r}"
  assert False  # unreachable

_canonicalize_partition = _jax.canonicalize_partition  # type: ignore
_canonicalize_partitions = _jax.canonicalize_partitions  # type: ignore

def _get_default_unconstrained(): return _UNCONSTRAINED_PARTITION

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove the axis from the partitions entry (use None in that position) or drop it from reduced
2. Validate spec construction: assert set(flat_partitions).isdisjoint(reduced)

Example fix

# before
PartitionSpec(('data', 'model'), reduced=('data',))
# after
PartitionSpec((None, 'model'), reduced=('data',))

flat = {p for part in partitions for p in (part if isinstance(part, tuple) else (part,))}
assert flat.isdisjoint(reduced)

When it happens

Trigger: PartitionSpec(('data',), reduced=('data', 'model')) — 'data' is both a partitioning axis and a reduction axis.

Common situations: Spec generation that marks every mesh axis as reduced while also using them in partitions; migration from older sharding APIs.

Related errors

• unreduced cannot contain None. All elements in unreduced sho
• reduced cannot contain None. All elements in reduced should
• `unreduced` and `reduced` argument to PartitionSpec cannot o
• partitions cannot overlap with unreduced axes passed to Part
• Mapped away dimension of inputs passed to vmap should be sha

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