jax-ml/jax · error · ValueError

Mesh requires the ndim of its first argument (`devices`) to

Error message

Mesh requires the ndim of its first argument (`devices`) to equal the length of its second argument (`axis_names`), but got devices.ndim == {devices.ndim} and len(axis_names) == {len(axis_names)}.

What it means

Mesh requires devices.ndim == len(axis_names): each mesh axis corresponds to one named dimension of the devices ndarray. If you pass a 2D device array with 3 names (or forget to reshape a flat device list), ValueError is raised.

Source

Thrown at jax/_src/mesh.py:277

    obj = object.__new__(Mesh)
    object.__setattr__(obj, 'devices', devices)
    object.__setattr__(obj, 'axis_names', axis_names)
    object.__setattr__(obj, 'axis_types', axis_types)
    object.__setattr__(obj, 'size', size)
    return obj

  def __new__(cls, devices: np.ndarray | Sequence[xc.Device],
              axis_names: str | Sequence[MeshAxisName],
              axis_types: tuple[AxisType, ...] | None = None):
    if not isinstance(devices, np.ndarray):
      devices = np.array(devices)
    if isinstance(axis_names, str):
      axis_names = (axis_names,)
    axis_names = tuple(axis_names)
    if any(i is None for i in axis_names):
      raise ValueError(f"Mesh axis names cannot be None. Got: {axis_names}")
    if devices.ndim != len(axis_names):
      raise ValueError(
          "Mesh requires the ndim of its first argument (`devices`) to equal "
          "the length of its second argument (`axis_names`), but got "
          f"devices.ndim == {devices.ndim} and "
          f"len(axis_names) == {len(axis_names)}.")

    devices_flat = tuple(devices.flat)
    axis_types = _normalize_axis_types(axis_names, axis_types, 'Mesh',
                                       AxisType.Auto)
    empty = not axis_names and devices_flat[0] is None
    size = 0 if empty else math.prod(devices.shape)
    return cls._create(devices_flat, devices.shape, axis_names,
                       axis_types, size)

  # No __eq__ or __hash__: interned classes use object identity.

  @property
  def is_scalar(self):
    return self.size == 1 and not self.axis_names

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Match the reshape to the names: np.array(devices).reshape(x, y) with ('x','y') names
2. Wrap flat devices with an explicit reshape — Mesh never reshapes for you
3. Compute the shape programmatically: divide jax.device_count() by known axis sizes
4. Use create_device_mesh to build the device array with the intended shape

Example fix

# before
mesh = jax.sharding.Mesh(jax.devices(), ('x', 'y'))

# after
import numpy as np
mesh = jax.sharding.Mesh(np.array(jax.devices()).reshape(4, 8), ('x', 'y'))

import numpy as np
arr = np.array(devices)
assert arr.ndim == len(axis_names), f'{arr.ndim} != {len(axis_names)}'

When it happens

Trigger: Mesh(jax.devices(), ('x','y')) — 1D list with two names; devices reshaped to (4,4,2) but only ('data','model') given; forgetting that a raw device list is never reshaped automatically.

Common situations: The single most common Mesh construction error: writing a mesh shape in the names that doesn't match the reshape applied to devices; copy-pasting mesh definitions between machines with different device counts without updating the reshape.

Related errors

• Number of axis names should match the number of axis_types.
• Mesh axis names cannot be None. Got: {axis_names}
• 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/212668c1c39e481a. Report an issue: GitHub.