{"record":{"id":"d9ea639e0a72996e","repo":"jax-ml/jax","slug":"transpose-cannot-be-moved-before-a-tiling-transfor","errorCode":null,"errorMessage":"Transpose cannot be moved before a tiling transform when it changes the set of tiled dimensions. (permutation: {perm}, tiling: {self.tiling})","messagePattern":"Transpose cannot be moved before a tiling transform when it changes the set of tiled dimensions\\. \\(permutation: (.+?), tiling: (.+?)\\)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":786,"sourceCode":"      case state_types.AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case _:\n        raise TypeError(f\"Cannot transform type: {x}\")\n\n  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:\n    return TilingTransform(self.tiling)\n\n  def commute_transpose(\n      self, _: jax_core.AbstractValue,\n      transpose: state_types.TransposeTransform,\n  ) -> tuple[state_types.TransposeTransform, UntilingTransform]:\n    # The transpose in question is applied to the untiled ref so we\n    # need to translate it by duplicating and offsetting the last part.\n    perm = transpose.permutation\n    off = len(perm)\n    new_suffix = [i + off for i in perm[-len(self.tiling) :]]\n    if set(new_suffix) != set(range(off, off + len(self.tiling))):\n      raise ValueError(\n          \"Transpose cannot be moved before a tiling transform when it changes\"\n          f\" the set of tiled dimensions. (permutation: {perm}, tiling:\"\n          f\" {self.tiling})\"\n      )\n\n    new_tiling = tuple(self.tiling[i - off] for i in new_suffix)\n    new_transpose = state_types.TransposeTransform((*perm, *new_suffix))\n    return new_transpose, dataclasses.replace(self, tiling=new_tiling)\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, UntilingTransform]:\n    del aval\n    idxs = indexer.indices\n    indexer_shape = indexer.shape\n    untiled_idxs = idxs[: -len(self.tiling)]\n    tiled_idxs = idxs[-len(self.tiling) :]\n    idxs_after_tiling: list[indexing.Slice] = []","sourceCodeStart":768,"sourceCodeEnd":804,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L768-L804","documentation":"When a transpose transform is commuted past a TilingTransform, the transpose must permute only the non-tiled leading dimensions — it cannot change which dimensions are tiled. If the permutation, offset by the untiled rank, does not map onto the tiled suffix of dimensions, ValueError is raised because the tiling could not be consistently re-derived.","triggerScenarios":"Applying a transpose on a tiled ref that moves a tiled dimension into a leading position (or vice versa) and then triggering transform commutation during lowering, e.g. transposing a 2D tiled block spec before passing to a kernel.","commonSituations":"Transposing blocked matrices in Pallas kernels where the tile layout is fixed (e.g. swapping rows/cols of a (128, N)-tiled TMEM accumulator); chained transforms (tiling then transpose) produced by automatic transform compositions in newer JAX versions.","solutions":["Transpose before tiling, i.e. apply the transpose to the untiled ref so the tiled dimensions stay last","Restructure the kernel to index the transposed view explicitly instead of commuting transforms","Choose a tiling whose tiled dimensions are invariant under your permutation (permute only leading dims)"],"exampleFix":"# before\nref_t = tiling_transform.apply(ref)\nref_tt = transpose_transform.apply(ref_t)  # ValueError if tiled dims move\n\n# after\nref_t = transpose_transform.apply(ref)   # transpose first, on untiled ref\nref_tt = tiling_transform.apply(ref_t)    # tiling applied after","handlingStrategy":"validation","validationCode":"# ensure the transpose only permutes leading (untiled) dims\nn_tiled = len(tiling)\nassert set(perm[-n_tiled:]) == set(\n    range(len(perm) - n_tiled, len(perm)\n)), 'transpose moves tiled dimensions'","typeGuard":"def transpose_preserves_tiled_dims(perm, n_tiled) -> bool:\n    return sorted(perm[-n_tiled:]) == list(\n        range(len(perm) - n_tiled, len(perm))\n    )","tryCatchPattern":"null","preventionTips":["Apply transposes before tiling transforms","Keep tiled dimensions as the trailing axes"],"tags":["jax","pallas","mosaic-gpu","transpose","tiling","transforms"],"backgroundTag":"transform-commutation-failed","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}