jax-ml/jax · error · ValueError

Offset {i} is not divisible by tile size {t}

Error message

Offset {i} is not divisible by tile size {t}

What it means

Tiling requires offsets to be aligned to tile boundaries; an offset not divisible by the corresponding tile size cannot be expressed in the tiled index space.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:1006

  return tuple(gmem_transforms)


def tile_offset(
    offsets: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing offsets in `offsets` according to `tiling`.

  Raises if the offsets are not aligned with the start of a tile.
  """
  if len(offsets) < len(tiling):
    raise ValueError(f"Offsets {offsets} have lower rank than tiling {tiling}")
  untiled_offsets, tiled_offsets = (
      offsets[: -len(tiling)],
      offsets[-len(tiling) :],
  )
  for i, t in zip(tiled_offsets, tiling, strict=True):
    if i % t != 0:
      raise ValueError(f"Offset {i} is not divisible by tile size {t}")
  return (
      *untiled_offsets,
      *[i // t for i, t in zip(tiled_offsets, tiling, strict=True)],
      *[0] * len(tiling),
  )


def tile_strides(
    strides: tuple[int, ...], tiling: tuple[int, ...]
) -> tuple[int, ...]:
  """Tiles the trailing strides in `strides` according to `tiling`.

  The `len(tiling)` trailing strides in `strides` must be the `len(tiling)`
  smallest strides in `strides`. The same property holds in the result, i.e.,
  given two tiles with indices i and j (i < j) with strides tiled according to
  this function, then all the elements in tile i are physically ordered before
  all the elements in tile j.

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

  1. Round offsets down/up to the nearest multiple of the tile size
  2. Adjust the tile size so offsets align
  3. Do the fine-grained offsetting with arithmetic on the loaded tile instead

Example fix

// before
tile_offset((0, 5), (8, 8))
// after
tile_offset((0, 8), (8, 8))  # aligned; handle +5 within the tile
Defensive patterns

Strategy: validation

Validate before calling

assert all(o % t == 0 for o, t in zip(offsets[-len(tiling):], tiling)), 'offsets must be tile-aligned'

Prevention

When it happens

Trigger: tile_offset where any trailing offset % tile_size != 0, e.g. offset 5 with tile size 8.

Common situations: Slicing a tiled buffer at a non-tile-aligned offset (e.g. slicing at element 16 in tiles of 32).

Related errors


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