jax-ml/jax · error · TypeError

axis_types passed to {name} must be of type `jax.sharding.Ax

Error message

axis_types passed to {name} must be of type `jax.sharding.AxisType`. Got {axis_types} of type {tuple(type(a) for a in axis_types)}

What it means

Mesh/AbstractMesh __new__ normalizes the axis_types argument (defaulting to a tuple of AxisType.AUTO) and requires every element to be a jax.sharding.AxisType enum member. Passing strings like 'manual' or arbitrary objects raises TypeError via _normalize_axis_types.

Source

Thrown at jax/_src/mesh.py:126

  return Mesh(global_mesh.devices[subcube_indices_tuple], global_mesh.axis_names)


class AxisType(enum.Enum):
  Auto = enum.auto()
  Explicit = enum.auto()
  Manual = enum.auto()

  def __repr__(self):
    return self.name

def _normalize_axis_types(axis_names, axis_types, name, default_axis_type):
  axis_types = ((default_axis_type,) * len(axis_names)
                if axis_types is None else axis_types)
  if not isinstance(axis_types, tuple):
    axis_types = (axis_types,)

  if not all(isinstance(a, AxisType) for a in axis_types):
    raise TypeError(
        f"axis_types passed to {name} must be of type `jax.sharding.AxisType`."
        f" Got {axis_types} of type {tuple(type(a) for a in axis_types)}")
  if len(axis_names) != len(axis_types):
    raise ValueError(
        "Number of axis names should match the number of axis_types. Got"
        f" axis_names={axis_names} and axis_types={axis_types}")
  return axis_types

def all_axis_types_match(axis_types, ty: AxisType) -> bool:
  if not axis_types:
    return False
  return all(t == ty for t in axis_types)

def any_axis_types_match(axis_types, ty: AxisType) -> bool:
  if not axis_types:
    return False
  return any(t == ty for t in axis_types)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use jax.sharding.AxisType members: AxisType.Auto, AxisType.Manual, AxisType.Explicit
2. Pass a tuple of length len(axis_names), a single AxisType (broadcast), or None
3. Check spelling of the enum attribute against the installed JAX version's jax.sharding.AxisType

Example fix

# before
mesh = jax.sharding.Mesh(devs, ('data', 'model'), axis_types=('auto', 'manual'))

# after
from jax.sharding import AxisType
mesh = jax.sharding.Mesh(devs, ('data', 'model'), axis_types=(AxisType.Auto, AxisType.Manual))

from jax.sharding import AxisType
ok = axis_types is None or all(isinstance(a, AxisType) for a in
    (axis_types if isinstance(axis_types, tuple) else (axis_types,)))

from jax.sharding import AxisType
def valid_axis_types(t):
    if t is None: return True
    ts = t if isinstance(t, tuple) else (t,)
    return all(isinstance(a, AxisType) for a in ts)

try:
    mesh = jax.sharding.Mesh(devs, names, axis_types=axis_types)
except TypeError as e:
    if 'AxisType' in str(e):
        mesh = jax.sharding.Mesh(devs, names)  # default Auto
    else: raise

When it happens

Trigger: Constructing jax.sharding.Mesh(devices, axis_names, axis_types=['manual']) or passing a single non-AxisType scalar; converting code that used string sharding types from other frameworks.

Common situations: Newer JAX versions adding the AxisType (auto/manual/explicit) parameter; translating PyTorch/FSDP-style 'manual' string flags; LLM-training configs that specify mesh autosharding mode.

Understand the failure class

Background: Invalid enum value errors: "Unknown type", "Invalid scope", "must be one of" — when a string is not on the library's allowed list — this error's family across 23 libraries.

Related errors

• Number of axis names should match the number of axis_types.
• Expected mesh of type `jax.sharding.AbstractMesh`. Got type:
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• callbacks are only supported in spmd computations when all m

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