{"record":{"id":"b433570a7f212eb1","repo":"jax-ml/jax","slug":"tiling-and-grid-must-have-same-length","errorCode":null,"errorMessage":"{tiling=} and {grid=} must have same length.","messagePattern":"(.+?) and (.+?) must have same length\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/helpers.py","lineNumber":131,"sourceCode":"  +------------+------------------+\n  |     2      | (0, 2)           |\n  +------------+------------------+\n  |     3      | (1, 0)           |\n  +------------+------------------+\n\n  If ``init_carry`` is passed then ``nd_loop()`` will expect the body to\n  take and return the carry. If it's ``None`` then no carry argument is\n  expected.\n\n  See also:\n    - :func:`jax.experimental.pallas.loop`: A loop over a single dimension.\n  \"\"\"\n\n  axis_index = lax.axis_index(collective_axes)\n  axis_size = lax.axis_size(collective_axes)\n  if tiling:\n    if len(grid) != len(tiling):\n      raise ValueError(f\"{tiling=} and {grid=} must have same length.\")\n    for dim, tile in zip(grid, tiling, strict=True):\n      if isinstance(dim, (int, np.integer)) and dim % tile != 0:\n        raise ValueError(f\"Tiling {tiling} does not divide grid {grid}.\")\n    tile_grid = tuple(\n        dim // tile for dim, tile in zip(grid, tiling, strict=True))\n    grid = (*tile_grid, *tiling)\n\n  grid_size = 1\n  for dim in grid:\n    grid_size = grid_size * dim\n  grid_size = jnp.astype(grid_size, axis_index.dtype)\n\n  def decorator(body):\n    def wrapper(wave_step, carry):\n      nonlocal body\n      step = wave_step * axis_size + axis_index\n      # The loop below is conceptually ``jnp.unravel_index``, but it uses\n      # ``lax`` APIs instead of ``jax.numpy`` to minimize the number of","sourceCodeStart":113,"sourceCodeEnd":149,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/helpers.py#L113-L149","documentation":"Raised by pallas.mosaic_gpu.helpers.nd_loop when a `tiling` tuple is provided whose length differs from `grid`. Tiling pairs each grid dimension with a tile size, so mismatched lengths make the tiling undefined.","triggerScenarios":"Calling nd_loop(grid=(128, 128), tiling=(64,)) — any call where len(tiling) != len(grid).","commonSituations":"Adding a grid dimension for a new kernel axis but forgetting to extend tiling; reusing a tiling constant across kernels with different dimensionality; computing tiling from block shapes that don't cover all grid dims.","solutions":["Give tiling one entry per grid dim: nd_loop(grid=(128, 128), tiling=(64, 64))","Derive tiling programmatically: tiling=tuple(min(t, g) for t, g in zip(tiles, grid)) and assert lengths up front"],"exampleFix":"// before\nnd_loop(grid=(128, 128), tiling=(64,))\n\n// after\nnd_loop(grid=(128, 128), tiling=(64, 64))","handlingStrategy":"validation","validationCode":"assert len(grid) == len(tiling), f'{grid=} {tiling=}'","typeGuard":"def tiling_valid(grid, tiling) -> bool:\n    return len(tiling) == len(grid)","tryCatchPattern":"try:\n    nd_loop(..., grid=grid, tiling=tiling)\nexcept ValueError:\n    tiling = tiling[:len(grid)] + (1,) * (len(grid) - len(tiling))\n    nd_loop(..., grid=grid, tiling=tiling)","preventionTips":["Define tiling next to grid and derive lengths together","Add an assert on lengths before calling nd_loop"],"tags":["jax","pallas","mosaic-gpu","nd-loop","tiling","shape-mismatch"],"backgroundTag":"config-length-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}