jax-ml/jax · error · ValueError

`unreduced_kind` should be `None` when `unreduced` is an emp

Error message

`unreduced_kind` should be `None` when `unreduced` is an empty set. Got {unreduced_kind=} and {unreduced=}

What it means

unreduced_kind is only meaningful when the unreduced set is non-empty. Passing a non-None unreduced_kind together with an empty unreduced set is contradictory and raises ValueError.

Source

Thrown at jax/_src/partition_spec.py:42

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Set unreduced_kind=None when unreduced is empty: kind = kind if unreduced else None
2. Default the kind argument to None and only set it when axes exist

Example fix

# before
PartitionSpec(ps, unreduced=(), unreduced_kind=kind)
# after
PartitionSpec(ps, unreduced=(), unreduced_kind=kind if axes else None)

unreduced_kind = unreduced_kind if unreduced else None

When it happens

Trigger: PartitionSpec(ps, unreduced=(), unreduced_kind=UnreducedKind.PARTIAL); often happens when the unreduced set is computed dynamically and ends up empty while kind is hard-coded.

Common situations: Config-driven code that always sets a kind but sometimes has no unreduced axes.

Related errors

• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• Axes mentioned in `manual_axis_type` field of ShapedArray sh
• varying and unreduced cannot have common mesh axes. Got vary
• varying and reduced cannot have common mesh axes. Got varyin
• `unreduced_kind` should be `None` when `unreduced` is an emp

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