jax-ml/jax · error · ValueError

partitions cannot overlap with unreduced axes passed to…

Error message

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

What it means

In the new-style PartitionSpec, axes listed in the positional partitions (classic sharding mapping) cannot also appear in the unreduced set — an axis is either a partitioning axis or an unreduced axis, not both.

Solutions

Remove the axis from partitions or from unreduced so each appears in only one
If the axis should stay un-reduced, leave it out of the positional partitions entirely

Example fix

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

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Each mesh axis should appear in exactly one role in a spec

When it happens

Trigger: PartitionSpec(('data',), unreduced=('data',)) — the same mesh axis used both as a partition spec entry and in unreduced.

Common situations: Reusing an axis name in a spec template; generated specs where partitions and unreduced are filled from the same axis pool.

Related errors

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

Appendix: source

Thrown at jax/_src/partition_spec.py:50

        " to the mesh axes.")
  if unreduced & reduced:
    raise ValueError(
        "`unreduced` and `reduced` argument to PartitionSpec cannot overlap. "
        f"Got unreduced: {unreduced} and reduced: {reduced}")
  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

View on GitHub (pinned to 1e1c6a8fc0)