{"record":{"id":"7d3bea311b1614eb","repo":"jax-ml/jax","slug":"tiling-tiling-does-not-divide-grid-grid","errorCode":null,"errorMessage":"Tiling {tiling} does not divide grid {grid}.","messagePattern":"Tiling (.+?) does not divide grid (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/helpers.py","lineNumber":134,"sourceCode":"  |     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\n      # primitives used.\n      index = []\n      for grid_dim in reversed(grid):","sourceCodeStart":116,"sourceCodeEnd":152,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/helpers.py#L116-L152","documentation":"nd_loop validates that each concrete (integer) grid dimension is divisible by its corresponding tile size, since the loop grid becomes (*tile_grid, *tiling). Non-divisible pairs would create partial tiles and are rejected.","triggerScenarios":"nd_loop(grid=(100,), tiling=(64,)) — 100 % 64 != 0. Only applies when the grid dim is an int/np.integer (dynamic tracer dims are skipped).","commonSituations":"Hardcoding tile sizes (e.g. 64) while problem sizes vary (e.g. M=100); changing block sizes in a matmul kernel without padding the grid to a tile multiple.","solutions":["Pad the grid up to a multiple of the tile: grid=(128,), tiling=(64,)","Or choose a tile that divides the grid: tiling=(50,) or (25,) for grid=(100,)"],"exampleFix":"// before\nnd_loop(grid=(100,), tiling=(64,))\n\n// after\ngrid = (-(-m // 64) * 64,)  # ceil to multiple of 64\nnd_loop(grid=grid, tiling=(64,))","handlingStrategy":"validation","validationCode":"import numbers\nfor d, t in zip(grid, tiling, strict=True):\n    if isinstance(d, numbers.Integral) and d % t:\n        raise ValueError(f'{d} not divisible by tile {t}')","typeGuard":"def grid_divides_tiling(grid, tiling) -> bool:\n    return all(not isinstance(d, int) or d % t == 0 for d, t in zip(grid, tiling))","tryCatchPattern":"try:\n    nd_loop(..., grid=grid, tiling=tiling)\nexcept ValueError:\n    grid = tuple(-(-d // t) * t for d, t in zip(grid, tiling))  # ceil-padded\n    nd_loop(..., grid=grid, tiling=tiling)","preventionTips":["Ceil-pad problem sizes to tile multiples before launching","Prefer power-of-two tiles that divide padded sizes"],"tags":["jax","pallas","mosaic-gpu","nd-loop","tiling","divisibility"],"backgroundTag":"size-not-divisible-by-tile","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}