jax-ml/jax · error · ValueError

Mesh axis names cannot be None. Got: {axis_names}

Error message

Mesh axis names cannot be None. Got: {axis_names}

What it means

Mesh._create converts axis_names to a tuple and rejects any None entries with ValueError. Mesh axis names are used as dict keys for shardings and resource lookups, so None names would break every downstream mapping.

Source

Thrown at jax/_src/mesh.py:275

    devices = np.array(flat_devices_tuple).reshape(device_shape)
    devices.flags.writeable = False
    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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Filter out None axis names and reshape devices accordingly
2. Default missing axis names to a real name like 'singleton' or omit that axis entirely
3. Validate config-derived names before constructing the Mesh

Example fix

# before
names = ('data', maybe_tensor_name)  # maybe_tensor_name may be None
mesh = jax.sharding.Mesh(devs.reshape(8, 1), names)

# after
names = tuple(n for n in (base, maybe_tensor_name) if n)
mesh = jax.sharding.Mesh(devs.reshape([d for d in devs.shape if d != 1][:len(names)] or (devs.size,)), names)

names = tuple(n for n in candidate_names if n is not None)
assert all(names), 'no None axis names'

def valid_axis_names(names):
    return all(n is not None for n in names)

When it happens

Trigger: Passing axis_names containing None, e.g. Mesh(devs, ('data', None)) or a list built by zipping mismatched sequences; also Mesh(devs, (None,)) when a name variable failed to be set.

Common situations: Programmatic construction of mesh names from config where an optional axis (e.g. 'tensor' absent for some runs) yields None; defaulting missing config entries to None instead of skipping them.

Related errors

• Cannot update the mesh of the current resource environment.
• Mesh requires the ndim of its first argument (`devices`) to
• jax.shard_map requires axis_names={axis_names} to be a subse
• Mapped away dimension of inputs passed to vmap should be sha
• Unmapped values passed to vmap cannot be sharded along the m

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