{"record":{"id":"5095e78413ebf0d4","repo":"jax-ml/jax","slug":"cannot-commute-untilingtransform-with-reshapetr","errorCode":null,"errorMessage":"Cannot commute `UntilingTransform` with `ReshapeTransform` when any of the dimensions being collapsed other than the minormost one has non-unit tiling. Attempted to reshape tiled slice of shape {before} into shape {after}","messagePattern":"Cannot commute `UntilingTransform` with `ReshapeTransform` when any of the dimensions being collapsed other than the minormost one has non-unit tiling\\. Attempted to reshape tiled slice of shape (.+?) into shape (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":917,"sourceCode":"      rev_tiling_slice = rev_tiling_to_process[:ndim]\n      rev_shape_slice = rev_shape_to_process[:ndim]\n      new_tiling_dim = math.prod(rev_tiling_slice)\n      num_elems = math.prod(rev_shape_slice)\n      assert num_elems % new_tiling_dim == 0\n      new_tiled_dim = num_elems // new_tiling_dim\n      # If any other dimension than the minormost one has non-unit tiling, then\n      # we cannot commute the reshape and untile transforms.\n      #\n      # Note that we could also support the case where we are collapsing\n      # trailing tiled dimensions where the tile size is the dimension size\n      # (i.e. there is a single tile).\n      if any(t != 1 for t in rev_tiling_slice[1:]):\n        before = (\n            *[s // t for s, t in zip(rev_shape_slice, rev_tiling_slice)],\n            *rev_tiling_slice,\n        )\n        after = (new_tiled_dim, new_tiling_dim)\n        raise ValueError(\n            \"Cannot commute `UntilingTransform` with `ReshapeTransform` when \"\n            \"any of the dimensions being collapsed other than the minormost \"\n            \"one has non-unit tiling. Attempted to reshape tiled slice of \"\n            f\"shape {before} into shape {after}\"\n        )\n      rev_new_tiling.append(new_tiling_dim)\n      rev_new_tiled_dims.append(new_tiled_dim)\n      rev_tiling_to_process = rev_tiling_to_process[ndim:]\n      rev_shape_to_process = rev_shape_to_process[ndim:]\n      if not rev_tiling_to_process:\n        break\n    assert not rev_tiling_to_process\n    assert not rev_shape_to_process\n    new_tiling = tuple(rev_new_tiling[::-1])\n    new_tiled_dims = tuple(rev_new_tiled_dims[::-1])\n    new_shape = (\n        *transform.shape[:len(components) - len(rev_new_tiling)],\n        *new_tiled_dims,","sourceCodeStart":899,"sourceCodeEnd":935,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L899-L935","documentation":"When collapsing several dimensions into one during untile/reshape commutation, every collapsed dimension except the minormost must have unit tiling. If a non-minor collapsed dim is tiled, the memory layout cannot be preserved and a ValueError is raised with the offending before/after shapes.","triggerScenarios":"Reshaping a block whose tiled slice has non-unit tiling on a non-minormost collapsed dimension, e.g. shape (2,2) tiled (2,2) reshaped to (4,) / tiling (4).","commonSituations":"Swizzled or tiled layouts (common with TMA block specs with tiling like (8, 128)) combined with a flatten/reshape across the major dimension.","solutions":["Adjust the BlockSpec tiling so non-minormost dims of the collapsed group have tiling 1","Reshape only within one tiled dimension rather than across tiled dimensions","Materialize the block (copy out of the transformed ref) before the aggressive reshape"],"exampleFix":"// before\n# tiling (2, 2), reshape (2,2)->(4,)\nblock = block.reshape((4,))\n\n// after\n# use tiling (1, 2) so only the minor dim is tiled\nblock = block.reshape((4,))","handlingStrategy":"validation","validationCode":"# before reshape, check non-minor collapsed dims have unit tiling\nassert all(t == 1 for t in tiling[:-1]), 'non-minor dims must be untiled'","typeGuard":null,"tryCatchPattern":"try: reshape\\nexcept ValueError as e: adjust BlockSpec tiling","preventionTips":["Set BlockSpec tiling so only the minormost dim is tiled when flattening"],"tags":["jax","pallas","mosaic-gpu","reshape","tiling","layout"],"backgroundTag":"unsupported-transform-commutation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}