{"record":{"id":"de6af008d8497a8b","repo":"jax-ml/jax","slug":"grid-names-must-have-the-same-length-as-grid-got","errorCode":null,"errorMessage":"grid_names must have the same length as grid, got {self}.","messagePattern":"grid_names must have the same length as grid, got (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1642,"sourceCode":"\n_WARPGROUP_AXIS_NAME = object()\n\n@dataclasses.dataclass(frozen=True, kw_only=True)\nclass Mesh(pallas_core.Mesh):\n  grid: Sequence[int] = ()\n  grid_names: Sequence[str] = ()\n  cluster: Sequence[int] = ()\n  cluster_names: Sequence[str] = ()\n  # Those are NOT CUDA threads. On Hopper they correspond to warpgroups.\n  num_threads: int | None = None\n  thread_name: str | None = None\n  kernel_name: str | None = None\n\n  def __post_init__(self):\n    if len(self.cluster) > 3:\n      raise ValueError(f\"cluster= must be at most 3D, got {self}.\")\n    if len(self.grid_names) != len(self.grid):\n      raise ValueError(\n          f\"grid_names must have the same length as grid, got {self}.\"\n      )\n    if len(self.cluster_names) != len(self.cluster):\n      raise ValueError(\n          f\"cluster_names must have the same length as cluster, got {self}.\"\n      )\n    if (self.thread_name is None) != (self.num_threads is None):\n      raise ValueError(\n          \"num_threads and thread_name must be either both set or both None,\"\n          f\" got {self}\"\n      )\n    max_mosaic_threads = 2048 // 128\n    if self.num_threads is not None and self.num_threads > max_mosaic_threads:\n      raise ValueError(\n          \"Requested too many CUDA threads per block. Each Mosaic thread\"\n          f\" corresponds to 128 CUDA threads. At most {max_mosaic_threads}\"\n          f\" are supported, got {self}\"\n      )","sourceCodeStart":1624,"sourceCodeEnd":1660,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1624-L1660","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","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"],"exampleFix":"// before\ngrid=(8, 8)\ngrid_names=('i',)\n\n// after\ngrid=(8, 8)\ngrid_names=('i', 'j')","handlingStrategy":"validation","validationCode":"assert len(grid) == len(grid_names), (grid, grid_names)","typeGuard":"def names_match(names, dims) -> bool:\n    return len(names) == len(dims)","tryCatchPattern":"try:\n    cfg = LaunchConfig(grid=g, grid_names=names)\nexcept ValueError:\n    names = tuple(f'd{i}' for i in range(len(g)))\n    cfg = LaunchConfig(grid=g, grid_names=names)","preventionTips":["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"],"tags":["jax","pallas","mosaic-gpu","launch-config","validation","shape-mismatch"],"backgroundTag":"config-length-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}