jax-ml/jax · error · ValueError

`unreduced` and `reduced` argument to PartitionSpec cannot o

Error message

`unreduced` and `reduced` argument to PartitionSpec cannot overlap. Got unreduced: {unreduced} and reduced: {reduced}

What it means

A PartitionSpec axis cannot be simultaneously unreduced and reduced — the two sets are semantically exclusive (an axis either keeps its data un-reduced or gets reduced). Passing an axis in both sets raises this ValueError.

Source

Thrown at jax/_src/partition_spec.py:34

import enum
from typing import Any

from jax._src.util import weak_value_interner, immutable
from jax._src.lib import _jax

AxisName = Any

def _check(partitions, unreduced, reduced, unreduced_kind):
  if None in unreduced:
    raise ValueError(
        "unreduced cannot contain None. All elements in unreduced should refer"
        " to the mesh axes.")
  if None in reduced:
    raise ValueError(
        "reduced cannot contain None. All elements in reduced should refer"
        " 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:"

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Decide per axis whether it should be reduced or unreduced and remove it from the other set
2. Sanitize inputs: reduced = set(reduced) - set(unreduced) if unreduced takes precedence

Example fix

# before
PartitionSpec(('rep',), unreduced=('data',), reduced=('data', 'rep'))
# after
PartitionSpec(('rep',), unreduced=('data',), reduced=('rep',))

assert set(unreduced).isdisjoint(reduced), 'axis cannot be both unreduced and reduced'

When it happens

Trigger: PartitionSpec(partitions, unreduced={'data'}, reduced=('data',)) — typically from building both sets from the same list of mesh axes.

Common situations: Generated/config-driven sharding specs where the same axis name is included in both lists; merging specs programmatically.

Related errors

• unreduced cannot contain None. All elements in unreduced sho
• reduced cannot contain None. All elements in reduced should
• partitions cannot overlap with unreduced axes passed to Part
• partitions cannot overlap with reduced axes passed to Partit
• 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/88c38b9acb2fb224. Report an issue: GitHub.