jax-ml/jax · error · ValueError

grid_names must have the same length as grid, got {self}.

Error message

grid_names must have the same length as grid, got {self}.

What it means

Launch-config dataclass validation: `grid_names` (names for each grid axis) must have exactly the same length as `grid`. A mismatch means named-axis bookkeeping would be inconsistent, so construction fails immediately.

Source

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

_WARPGROUP_AXIS_NAME = object()

@dataclasses.dataclass(frozen=True, kw_only=True)
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}"
      )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

Make len(grid_names) == len(grid), e.g. grid=(8, 8), grid_names=('i', 'j')
If you don't need named axes, pass a default/empty naming consistent with the expected API

Example fix

// before
grid=(8, 8)
grid_names=('i',)

// after
grid=(8, 8)
grid_names=('i', 'j')

Defensive patterns

Strategy: validation

Validate before calling

assert len(grid) == len(grid_names), (grid, grid_names)

Type guard

def names_match(names, dims) -> bool:
    return len(names) == len(dims)

Try / catch

try:
    cfg = LaunchConfig(grid=g, grid_names=names)
except ValueError:
    names = tuple(f'd{i}' for i in range(len(g)))
    cfg = LaunchConfig(grid=g, grid_names=names)

Prevention

Generate axis names from the dims: tuple(f'd{i}' for i in range(len(grid)))
Validate config in one place (dataclass __post_init__ already does; mirror it in tests)
Add a unit test asserting config construction succeeds

When it happens

Trigger: Passing `grid=(8, 8)` with `grid_names=('i',)` (or any pair whose lengths differ) when constructing the kernel launch descriptor used by warp-specialized Mosaic GPU kernels.

Common situations: Adding or removing a grid dimension during development but forgetting to update grid_names; copy-pasting a config from a kernel with different dimensionality; using collective axes (e.g. from jax.lax.paxis) where names are auto-generated with different length.

Related errors

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