jax-ml/jax · error · TypeError

Expected unreduced_kind to be of type `jax.sharding.Unreduce

Error message

Expected unreduced_kind to be of type `jax.sharding.UnreducedKind` but got {type(unreduced_kind)}

What it means

The optional unreduced_kind argument of the new-style PartitionSpec must be a jax.sharding.UnreducedKind instance (or None). Passing any other type — string, int, custom class — raises TypeError.

Source

Thrown at jax/_src/partition_spec.py:38

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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jax.sharding.UnreducedKind values (import from jax.sharding) or omit the argument
2. Check type(unreduced_kind) is UnreducedKind before constructing

Example fix

# before
PartitionSpec(ps, unreduced=('d',), unreduced_kind='partial')
# after
from jax.sharding import UnreducedKind
PartitionSpec(ps, unreduced=('d',), unreduced_kind=UnreducedKind.PARTIAL)

from jax.sharding import UnreducedKind
assert unreduced_kind is None or isinstance(unreduced_kind, UnreducedKind)

def valid_kind(k): return k is None or isinstance(k, UnreducedKind)

When it happens

Trigger: PartitionSpec(..., unreduced=('data',), unreduced_kind='partial') or unreduced_kind=0.

Common situations: Passing a string flag from config (e.g. 'optimistic'/'partial') instead of the enum-like UnreducedKind object.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

• Length of sharding.spec ({len(out_s.spec)}) must be equal to
• Expected unreduced_kind to be of type `jax.sharding.Unreduce
• {name} argument of ManualAxisType should of type `frozenset`
• {kind} must be a tuple of factors
• rule must be a str, but got {type(rule)}

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