jax-ml/jax · error · ValueError

reduced cannot contain None. All elements in reduced should

Error message

reduced cannot contain None. All elements in reduced should refer to the mesh axes.

What it means

In the unreduced/reduced PartitionSpec API, the reduced set must contain only mesh axis names; None is rejected because a reduced axis always refers to a concrete mesh dimension that collectives will reduce over.

Source

Thrown at jax/_src/partition_spec.py:30

# See the License for the specific language governing permissions and
# limitations under the License.

from __future__ import annotations
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:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Filter None out: reduced = tuple(a for a in reduced if a is not None)
2. Only pass mesh axis names registered on your Mesh

Example fix

# before
reduced = (None, 'data')
# after
reduced = ('data',)

reduced = tuple(a for a in reduced if a is not None)

When it happens

Trigger: PartitionSpec(..., reduced=(None,)) or passing a spec built for classic sharding (with None) into the reduced argument.

Common situations: Migrating old PartitionSpec tuples containing None to the new API without filtering; programmatically building reduced sets from user input that includes None.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• unreduced cannot contain None. All elements in unreduced sho
• `unreduced` and `reduced` argument to PartitionSpec cannot o
• 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/6d104e92d9b3814d. Report an issue: GitHub.