jax-ml/jax · error · ValueError

Tiling {tiling} does not divide grid {grid}.

Error message

Tiling {tiling} does not divide grid {grid}.

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/helpers.py:134

  |     3      | (1, 0)           |
  +------------+------------------+

  If ``init_carry`` is passed then ``nd_loop()`` will expect the body to
  take and return the carry. If it's ``None`` then no carry argument is
  expected.

  See also:
    - :func:`jax.experimental.pallas.loop`: A loop over a single dimension.
  """

  axis_index = lax.axis_index(collective_axes)
  axis_size = lax.axis_size(collective_axes)
  if tiling:
    if len(grid) != len(tiling):
      raise ValueError(f"{tiling=} and {grid=} must have same length.")
    for dim, tile in zip(grid, tiling, strict=True):
      if isinstance(dim, (int, np.integer)) and dim % tile != 0:
        raise ValueError(f"Tiling {tiling} does not divide grid {grid}.")
    tile_grid = tuple(
        dim // tile for dim, tile in zip(grid, tiling, strict=True))
    grid = (*tile_grid, *tiling)

  grid_size = 1
  for dim in grid:
    grid_size = grid_size * dim
  grid_size = jnp.astype(grid_size, axis_index.dtype)

  def decorator(body):
    def wrapper(wave_step, carry):
      nonlocal body
      step = wave_step * axis_size + axis_index
      # The loop below is conceptually ``jnp.unravel_index``, but it uses
      # ``lax`` APIs instead of ``jax.numpy`` to minimize the number of
      # primitives used.
      index = []
      for grid_dim in reversed(grid):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the grid up to a multiple of the tile: grid=(128,), tiling=(64,)
2. Or choose a tile that divides the grid: tiling=(50,) or (25,) for grid=(100,)

Example fix

// before
nd_loop(grid=(100,), tiling=(64,))

// after
grid = (-(-m // 64) * 64,)  # ceil to multiple of 64
nd_loop(grid=grid, tiling=(64,))

import numbers
for d, t in zip(grid, tiling, strict=True):
    if isinstance(d, numbers.Integral) and d % t:
        raise ValueError(f'{d} not divisible by tile {t}')

def grid_divides_tiling(grid, tiling) -> bool:
    return all(not isinstance(d, int) or d % t == 0 for d, t in zip(grid, tiling))

try:
    nd_loop(..., grid=grid, tiling=tiling)
except ValueError:
    grid = tuple(-(-d // t) * t for d, t in zip(grid, tiling))  # ceil-padded
    nd_loop(..., grid=grid, tiling=tiling)

When it happens

Trigger: 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).

Common situations: 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.

Related errors

• {tiling=} and {grid=} must have same length.
• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Folding tiled dimensions into untiled dimensions is not supp
• Cannot commute `UntilingTransform` with `ReshapeTransform` w

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