jax-ml/jax · error · ValueError

{tiling=} and {grid=} must have same length.

Error message

{tiling=} and {grid=} must have same length.

What it means

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.

Source

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

  +------------+------------------+
  |     2      | (0, 2)           |
  +------------+------------------+
  |     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

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Give tiling one entry per grid dim: nd_loop(grid=(128, 128), tiling=(64, 64))
2. Derive tiling programmatically: tiling=tuple(min(t, g) for t, g in zip(tiles, grid)) and assert lengths up front

Example fix

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

// after
nd_loop(grid=(128, 128), tiling=(64, 64))

assert len(grid) == len(tiling), f'{grid=} {tiling=}'

def tiling_valid(grid, tiling) -> bool:
    return len(tiling) == len(grid)

try:
    nd_loop(..., grid=grid, tiling=tiling)
except ValueError:
    tiling = tiling[:len(grid)] + (1,) * (len(grid) - len(tiling))
    nd_loop(..., grid=grid, tiling=tiling)

When it happens

Trigger: Calling nd_loop(grid=(128, 128), tiling=(64,)) — any call where len(tiling) != len(grid).

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

Related errors

• Tiling {tiling} does not divide grid {grid}.
• 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/b433570a7f212eb1. Report an issue: GitHub.