jax-ml/jax · error · ValueError

Number of axis names should match the number of axis_types.

Error message

Number of axis names should match the number of axis_types. Got axis_names={axis_names} and axis_types={axis_types}

What it means

After type-checking, _normalize_axis_types verifies that the number of axis_types equals the number of axis_names, since each mesh axis needs exactly one type. A mismatch (e.g. 3 names, 2 types) raises ValueError with both values shown.

Source

Thrown at jax/_src/mesh.py:130

  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)


class BaseMesh:
  axis_names: tuple[MeshAxisName, ...]
  shape_tuple: tuple[tuple[str, int], ...]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Make len(axis_types) == len(axis_names) (or pass one AxisType scalar to broadcast)
2. Add AxisType.Auto for axes you don't want to specialize
3. Centralize names and types in one config dict so they stay in sync

Example fix

# before
mesh = jax.sharding.Mesh(devs, ('data','fsdp','model'),
                         axis_types=(AxisType.Auto, AxisType.Manual))

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

names = ('data','fsdp','model')
types = (AxisType.Auto, AxisType.Auto, AxisType.Manual)
assert len(types) == len(names) or not isinstance(types, tuple)

When it happens

Trigger: Mesh(devices, ('a','b','c'), axis_types=(AxisType.Auto, AxisType.Manual)) — length mismatch. Note a single non-tuple AxisType is broadcast, so this only fires for tuples of wrong length or a scalar string expanded wrongly.

Common situations: Adding/removing a mesh axis (e.g. inserting a 'fsdp' axis) without updating the parallel axis_types list in a training config; copy-paste between configs with different mesh shapes.

Related errors

• axis_types passed to {name} must be of type `jax.sharding.Ax
• Mesh requires the ndim of its first argument (`devices`) to
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m
• unexpected JAX type (e.g. shape/dtype) for argument to VJP f

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