jax-ml/jax · error · TypeError

Expected unreduced_kind to be of type…

Error message

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

What it means

Type guard in _check_mat, which validates the axis-type matrix (varying/unreduced/reduced) of an abstract value: the optional unreduced_kind discriminator must be an instance of jax.sharding.UnreducedKind (or None); anything else is a programming error in how the mat was constructed.

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.

Understand the failure class

Background: "invalid argument", "unknown mode", "not supported": invalid enum-like argument errors explained — this error's family across 19 libraries.

Related errors

• : arrays must have same number of dimensions, got .
• Expected unreduced_kind to be of type…
• factor_sizes must be a dict of str to int, but got
• must be a tuple of factors
• argument of ManualAxisType should of type `frozenset` or…

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

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