jax-ml/jax · error · ValueError

unreduced cannot contain None. All elements in unreduced sho

Error message

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

What it means

When constructing a PartitionSpec with unreduced axes (new-style sharding specs), every element of the unreduced set must be a real mesh axis name. None is not allowed because unreduced axes denote named mesh axes that remain un-reduced, unlike classic ParsedPartitionSpecs where None means 'not partitioned'.

Source

Thrown at jax/_src/partition_spec.py:26

#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# 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=}")

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Remove None entries from unreduced; only include actual mesh axis names
2. Use None only in the positional partitions part, not in unreduced/reduced

Example fix

# before
PartitionSpec(('data',), unreduced=(None, 'model'))
# after
PartitionSpec(('data',), unreduced=('model',))

assert None not in unreduced, 'unreduced takes mesh axis names only'

When it happens

Trigger: Creating a PartitionSpec with unreduced=(None, 'data') or reusing an old spec with None placeholders in the new unreduced API.

Common situations: Mixing classic NamedSharding-style specs (where None is idiomatic) with the newer unreduced/reduced PartitionSpec constructor; copy-pasting specs across APIs.

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

• reduced cannot contain None. All elements in reduced should
• `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/15f813a79bd14a5b. Report an issue: GitHub.