{"record":{"id":"21f8f3c17112dd99","repo":"jax-ml/jax","slug":"can-not-tile-strides-when-tiled-dimensions-have-be","errorCode":null,"errorMessage":"Can not tile strides when tiled dimensions have been transposed with untiled dimensions. Strides: {strides}, tiling: {tiling}","messagePattern":"Can not tile strides when tiled dimensions have been transposed with untiled dimensions\\. Strides: (.+?), tiling: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":1031,"sourceCode":"\ndef tile_strides(\n    strides: tuple[int, ...], tiling: tuple[int, ...]\n) -> tuple[int, ...]:\n  \"\"\"Tiles the trailing strides in `strides` according to `tiling`.\n\n  The `len(tiling)` trailing strides in `strides` must be the `len(tiling)`\n  smallest strides in `strides`. The same property holds in the result, i.e.,\n  given two tiles with indices i and j (i < j) with strides tiled according to\n  this function, then all the elements in tile i are physically ordered before\n  all the elements in tile j.\n\n  E.g., tile_strides((2048, 32, 1), (8, 4)) = (2048, 256, 32, 4, 1)\n  \"\"\"\n  if len(strides) < len(tiling):\n    raise ValueError(f\"Strides {strides} have lower rank than tiling {tiling}\")\n  ordered_strides = sorted(strides, reverse=True)\n  if set(ordered_strides[-len(tiling):]) != set(strides[-len(tiling):]):\n    raise ValueError(\n        \"Can not tile strides when tiled dimensions have been transposed with \"\n        f\"untiled dimensions. Strides: {strides}, tiling: {tiling}\"\n    )\n  untiled_strides, tiled_strides = strides[:-len(tiling)], strides[-len(tiling):]\n\n  # Zip the strides and tiling together, in order to sort them together. This\n  # allows handling cases where multiple tiling dimensions have the same stride,\n  # which can occur with size-1 dimensions.\n  tiled_strides_and_tiling: list[tuple[int, int]] = list(\n      zip(tiled_strides, tiling, strict=True))\n  tiled_ordered_strides_and_tiling = sorted(\n      tiled_strides_and_tiling, reverse=True)\n\n  to_ordered = lambda i: tiled_ordered_strides_and_tiling.index(tiled_strides_and_tiling[i])\n  from_ordered = lambda i: tiled_strides_and_tiling.index(tiled_ordered_strides_and_tiling[i])\n\n  ordered_tiling = [tiling[from_ordered(i)] for i in range(len(tiling))]\n  ordered_tiled_strides = [tiled_strides[from_ordered(i)] for i in range(len(tiling))]","sourceCodeStart":1013,"sourceCodeEnd":1049,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L1013-L1049","documentation":"Tiling assumes the trailing (tiled) dimensions carry the smallest strides; if the tiled dims have been transposed with untiled dims, the memory layout cannot be expressed as a simple tiling.","triggerScenarios":"tile_strides where the set of the len(tiling) smallest sorted strides differs from the set of the trailing strides, i.e. tiled dims are not the innermost-contiguous ones.","commonSituations":"Applying a tile transform to a transposed or column-major memref where inner strides don't correspond to the tiled dimensions.","solutions":["Transpose the data in memory first so tiled dims are innermost with smallest strides","Apply the tile transform before the transpose","Use tile sizes of 1 on the transposed dims to effectively skip tiling them"],"exampleFix":"// before\ntile_strides((1, 2048, 32), (8, 4))  # tiled dims transposed\n// after\nmem = transpose_to_row_major(mem)\ntile_strides((2048, 32, 1), (8, 4))","handlingStrategy":"validation","validationCode":"ordered = sorted(strides, reverse=True)\nassert set(ordered[-len(tiling):]) == set(strides[-len(tiling):]), 'tiled dims must be innermost contiguous'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep tiled dimensions innermost (row-major) before applying tile transforms","Transpose in memory, not via strides"],"tags":["jax","mosaic","gpu","tiling","strides","transpose"],"backgroundTag":"invalid-memory-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}