jax-ml/jax · error · ValueError

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

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.

Solutions

Decide per axis whether it should be reduced or unreduced and remove it from the other set
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',))

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

Keep axis role assignments in one config source

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

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

Appendix: 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)