{"record":{"id":"9c8a633c12a0ce28","repo":"jax-ml/jax","slug":"non-divisible-tiling","errorCode":null,"errorMessage":"Non-divisible tiling:","messagePattern":"Non-divisible tiling:","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":1794,"sourceCode":"        num_chunks=num_chunks,\n        chunk_size=chunk_size,\n        base_offset=self.get_base(chunk) if chunk is not None else None,\n    )\n\n\ndef tile_shape(shape, tiling):\n  if len(tiling) > len(shape):\n    raise ValueError(\n        \"Expected tiling to be at most rank of shape. Got tiling:\"\n        f\" {tiling} (rank: {len(tiling)}) and shape {shape} (rank:\"\n        f\" {len(shape)}).\"\n    )\n  if not tiling:\n    return shape\n  tiling_rank = len(tiling)\n  for s, t in zip(shape[-tiling_rank:], tiling):\n    if s % t:\n      raise ValueError(\"Non-divisible tiling:\", shape, tiling)\n  return (\n      *shape[:-tiling_rank],\n      *(s // t for s, t in zip(shape[-tiling_rank:], tiling)),\n      *tiling,\n  )\n\n\ndef warp_tree_reduce(value, op, group_size):\n  \"\"\"Reduce a value across the warpgroup.\"\"\"\n  assert bytewidth(value.type) == 4\n  assert 32 % group_size == 0 and group_size <= 32\n  i32 = ir.IntegerType.get_signless(32)\n  result = value\n  iters = np.log2(group_size)\n  if not iters.is_integer():\n    raise ValueError(\n        f\"Warp reduction group size should be a power of 2 (got {group_size})\"\n    )","sourceCodeStart":1776,"sourceCodeEnd":1812,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L1776-L1812","documentation":"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).","triggerScenarios":"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.","commonSituations":"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.","solutions":["Pad the shape dimension up to a multiple of the tiling factor (and slice results afterwards)","Choose a tiling factor that divides every trailing shape dimension","Restructure so the non-divisible leading dims are outside the tiled region (only trailing dims are tiled)"],"exampleFix":"# before\ntile_shape(shape=(1000, 128), tiling=(128, 128))\n# after\npadded = 1024  # next multiple of 128\ntile_shape(shape=(padded, 128), tiling=(128, 128))","handlingStrategy":"validation","validationCode":"import math\nassert all(s % t == 0 for s, t in zip(shape[-len(tiling):], tiling)), 'non-divisible tiling'\nresult = tile_shape(shape, tiling)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Pad sequence/feature dims to multiples of the block size before kernel launch","Parametrize tests over ragged dims to catch divisibility breaks early"],"tags":["mosaic-gpu","tiling","divisibility"],"backgroundTag":"non-divisible-shape-tiling","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}