jax-ml/jax · error · ValueError

unreduced cannot contain None. All elements in unreduced…

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'.

Solutions

Remove None entries from unreduced; only include actual mesh axis names
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',))

Defensive patterns

Strategy: validation

Validate before calling

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

Prevention

None belongs only in positional partitions, never unreduced/reduced

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

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

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