jax-ml/jax · error · TypeError

{name} argument of ManualAxisType should of type `frozenset`

Error message

{name} argument of ManualAxisType should of type `frozenset` or `set`. Got type {type(val)}

What it means

_canonicalize_mat accepts only frozenset or set for ManualAxisType's varying/unreduced/reduced arguments. Passing a list, tuple, string, or other iterable raises TypeError.

Source

Thrown at jax/_src/core.py:2366

  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
class ManualAxisType:
  __slots__ = ('varying', 'unreduced', 'reduced', 'unreduced_kind',
               '__weakref__')

  varying: frozenset
  unreduced: frozenset
  reduced: frozenset
  unreduced_kind: UnreducedKind | None

  @staticmethod
  @weak_value_interner

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Wrap in set()/frozenset(): ManualAxisType(varying={'x','y'})
2. Validate inputs before constructing if they come from user config
3. Update any tuple-producing callers to convert

Example fix

// before
mat = ManualAxisType(varying=['x', 'y'])

// after
mat = ManualAxisType(varying={'x', 'y'})

varying = frozenset(varying) if not isinstance(varying, frozenset) else varying
# note: only set/frozenset accepted; convert lists first

def is_ok_axis_container(v): return isinstance(v, (set, frozenset))

When it happens

Trigger: ManualAxisType(varying=['x','y']) or varying=('x',) — lists/tuples are rejected rather than coerced.

Common situations: Naturally writing a list literal for the axis names; copy-pasting from reprs that show different containers; older code assuming lenient coercion.

Related errors

• Expected unreduced_kind to be of type `jax.sharding.Unreduce
• {kind} must be a tuple of factors
• rule must be a str, but got {type(rule)}
• factor_sizes must be a dict of str to int, but got {factor_s
• {}: arrays must have same number of dimensions, got {}.

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