jax-ml/jax · error · ValueError

cluster_names must have the same length as cluster, got {sel

Error message

cluster_names must have the same length as cluster, got {self}.

What it means

Launch-config dataclass validation: `cluster_names` must have the same length as `cluster`. Cluster axis names are used for lax.axis_index-style queries inside the kernel, so a mismatch is rejected at construction.

Source

Thrown at jax/_src/pallas/mosaic_gpu/core.py:1646

class Mesh(pallas_core.Mesh):
  grid: Sequence[int] = ()
  grid_names: Sequence[str] = ()
  cluster: Sequence[int] = ()
  cluster_names: Sequence[str] = ()
  # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.
  num_threads: int | None = None
  thread_name: str | None = None
  kernel_name: str | None = None

  def __post_init__(self):
    if len(self.cluster) > 3:
      raise ValueError(f"cluster= must be at most 3D, got {self}.")
    if len(self.grid_names) != len(self.grid):
      raise ValueError(
          f"grid_names must have the same length as grid, got {self}."
      )
    if len(self.cluster_names) != len(self.cluster):
      raise ValueError(
          f"cluster_names must have the same length as cluster, got {self}."
      )
    if (self.thread_name is None) != (self.num_threads is None):
      raise ValueError(
          "num_threads and thread_name must be either both set or both None,"
          f" got {self}"
      )
    max_mosaic_threads = 2048 // 128
    if self.num_threads is not None and self.num_threads > max_mosaic_threads:
      raise ValueError(
          "Requested too many CUDA threads per block. Each Mosaic thread"
          f" corresponds to 128 CUDA threads. At most {max_mosaic_threads}"
          f" are supported, got {self}"
      )
    object.__setattr__(self, "grid", tuple(self.grid))
    object.__setattr__(self, "grid_names", tuple(self.grid_names))
    object.__setattr__(self, "cluster", tuple(self.cluster))
    object.__setattr__(self, "cluster_names", tuple(self.cluster_names))

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Align lengths: cluster=(2, 2), cluster_names=('cq', 'ck')
Re-generate names with a comprehension: cluster_names=tuple(f'c{i}' for i in range(len(cluster)))

Example fix

// before
cluster=(2, 2)
cluster_names=('cq',)

// after
cluster=(2, 2)
cluster_names=('cq', 'ck')

Defensive patterns

Strategy: validation

Validate before calling

assert len(cluster) == len(cluster_names), (cluster, cluster_names)

Type guard

def cluster_names_valid(cluster, names) -> bool:
    return len(names) == len(cluster)

Try / catch

try:
    cfg = LaunchConfig(cluster=c, cluster_names=names)
except ValueError:
    names = tuple(f'c{i}' for i in range(len(c)))
    cfg = LaunchConfig(cluster=c, cluster_names=names)

Prevention

Derive cluster_names programmatically from cluster
Update names whenever cluster dimensionality changes
Test config construction in CI

When it happens

Trigger: Passing `cluster=(2, 2)` with `cluster_names=('cq',)` (lengths differ) to the Mosaic GPU launch/config dataclass.

Common situations: Changing cluster dimensionality (e.g. after hitting the 3D cluster limit) without updating names; deriving cluster_names from grid_names programmatically with an off-by-one or wrong source tuple.

Related errors

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