jax-ml/jax · error · ValueError

partitions cannot overlap with reduced axes passed to…

Error message

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

What it means

Validation in PartitionSpec._check: each leaf of the partitions specification names mesh axes to shard over, and an axis cannot simultaneously be sharded and pre-reduced. Naming an axis that also appears in the reduced set fails, with both sets reported.

Solutions

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

Example fix

# before
PartitionSpec(('data', 'model'), reduced=('data',))
# after
PartitionSpec((None, 'model'), reduced=('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(reduced)

Prevention

Use None in partition positions for axes that are 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

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

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