jax-ml/jax · error · ValueError

Non-divisible tiling:

Error message

Non-divisible tiling:

What it means

tile_shape requires each tiled dimension's size to be divisible by the corresponding tiling factor. Non-divisible tilings would produce fractional per-chunk extents, which the code cannot represent, so it raises ValueError('Non-divisible tiling:', shape, tiling).

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:1794

        num_chunks=num_chunks,
        chunk_size=chunk_size,
        base_offset=self.get_base(chunk) if chunk is not None else None,
    )


def tile_shape(shape, tiling):
  if len(tiling) > len(shape):
    raise ValueError(
        "Expected tiling to be at most rank of shape. Got tiling:"
        f" {tiling} (rank: {len(tiling)}) and shape {shape} (rank:"
        f" {len(shape)})."
    )
  if not tiling:
    return shape
  tiling_rank = len(tiling)
  for s, t in zip(shape[-tiling_rank:], tiling):
    if s % t:
      raise ValueError("Non-divisible tiling:", shape, tiling)
  return (
      *shape[:-tiling_rank],
      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),
      *tiling,
  )


def warp_tree_reduce(value, op, group_size):
  """Reduce a value across the warpgroup."""
  assert bytewidth(value.type) == 4
  assert 32 % group_size == 0 and group_size <= 32
  i32 = ir.IntegerType.get_signless(32)
  result = value
  iters = np.log2(group_size)
  if not iters.is_integer():
    raise ValueError(
        f"Warp reduction group size should be a power of 2 (got {group_size})"
    )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pad the shape dimension up to a multiple of the tiling factor (and slice results afterwards)
2. Choose a tiling factor that divides every trailing shape dimension
3. Restructure so the non-divisible leading dims are outside the tiled region (only trailing dims are tiled)

Example fix

# before
tile_shape(shape=(1000, 128), tiling=(128, 128))
# after
padded = 1024  # next multiple of 128
tile_shape(shape=(padded, 128), tiling=(128, 128))

import math
assert all(s % t == 0 for s, t in zip(shape[-len(tiling):], tiling)), 'non-divisible tiling'
result = tile_shape(shape, tiling)

When it happens

Trigger: Calling tile_shape with e.g. shape=(1000,) and tiling=(128,) since 1000 % 128 != 0. Happens when sequence lengths or vocab sizes are not multiples of the chosen block size.

Common situations: Using a power-of-two block size with a non-multiple tensor dimension (e.g. seq_len=1000, tile=128); switching a model config to a ragged/padded dimension without padding the tensor; changing tiling without re-checking dimension sizes.

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/9c8a633c12a0ce28. Report an issue: GitHub.