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

ManualAxisType's optional unreduced_kind annotation must be a jax.sharding.UnreducedKind (e.g. UnreducedKind.sum/min/max) or None. Passing a string, int, or other object is a type error.

Source

Thrown at jax/_src/core.py:2355

def get_memory_space(memory_space):
  assert memory_space is not None
  return memory_space


def _check_mat(varying, unreduced, reduced, unreduced_kind):
  if varying & unreduced:
    raise ValueError(
        "varying and unreduced cannot have common mesh axes. Got"
        f" varying={varying} and unreduced={unreduced}")
  if varying & reduced:
    raise ValueError(
        "varying and reduced cannot have common mesh axes. Got"
        f" varying={varying} and reduced={reduced}")
  assert not (varying & unreduced & 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=}")

def _canonicalize_mat(name, val):
  if not isinstance(val, frozenset):
    if not isinstance(val, set):
      raise TypeError(
          f"{name} argument of ManualAxisType should "
          f"of type `frozenset` or `set`. Got type {type(val)}")
    val = frozenset(val)
  return val


@immutable

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass jax.sharding.UnreducedKind.sum (or .min/.max) or None
2. On deserialization, map stored strings back to UnreducedKind members
3. Leave it None when there are no unreduced axes

Example fix

// before
mat = ManualAxisType(unreduced={'x'}, unreduced_kind='sum')

// after
from jax.sharding import UnreducedKind
mat = ManualAxisType(unreduced={'x'}, unreduced_kind=UnreducedKind.sum)

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: ManualAxisType(..., unreduced={'x'}, unreduced_kind='sum') or unreduced_kind=0 instead of the enum-like UnreducedKind instance.

Common situations: Assuming unreduced_kind is a free-form string; serializing/deserializing mats and losing the type; old JAX versions predating the typed enum.

Related errors

• `unreduced_kind` should be `None` when `unreduced` is an emp
• {name} argument of ManualAxisType should of type `frozenset`
• {name} cannot accept args with unreduced_kind={a.mat.unreduc
• {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/e30114756e033c12. Report an issue: GitHub.