{"record":{"id":"e4df649fbfee75df","repo":"jax-ml/jax","slug":"can-t-transpose-the-swizzled-dimension","errorCode":null,"errorMessage":"Can't transpose the swizzled dimension.","messagePattern":"Can't transpose the swizzled dimension\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1221,"sourceCode":"\n  def transform_type(self, x: jax_core.AbstractValue) -> jax_core.AbstractValue:\n    # Swizzling preserves the type\n    return x\n\n  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:\n    return SwizzleTransform(self.swizzle)\n\n  def swizzle_elems(self, dtype: jax.typing.DTypeLike | ir.Type) -> int:\n    if not isinstance(dtype, ir.Type):\n      dtype = mgpu_utils.dtype_to_ir_type(dtype)\n    return (self.swizzle * 8) // mgpu.bitwidth(dtype)\n\n  def commute_transpose(\n      self, _: jax_core.AbstractValue, transpose: state_types.TransposeTransform\n  ) -> tuple[state_types.TransposeTransform, UnswizzleRef]:\n    perm = transpose.permutation\n    if perm[-1] != len(perm) - 1:\n      raise ValueError(\"Can't transpose the swizzled dimension.\")\n    return transpose, self\n\n  def commute_reshape(\n      self, aval: jax_core.ShapedArray, transform: state_types.ReshapeTransform\n  ) -> tuple[state_types.ReshapeTransform, UnswizzleRef]:\n    shape = aval.shape\n    if shape[-1] != self.swizzle_elems(aval.dtype):\n      raise ValueError(\n          f\"Reshape shape {shape} is not divisible by swizzle elements\"\n          f\" {self.swizzle_elems(aval.dtype)}\"\n      )\n    return transform, self\n\n  def commute_ndindexer(\n      self, aval: jax_core.AbstractValue, indexer: indexing.NDIndexer\n  ) -> tuple[indexing.NDIndexer, UnswizzleRef]:\n    if not hasattr(aval, \"dtype\"):\n      raise ValueError(","sourceCodeStart":1203,"sourceCodeEnd":1239,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1203-L1239","documentation":"UnswizzleRef.commute_transpose raises when asked to commute a transpose past an unswizzle transform whose permutation moves the last (swizzled/minormost) dimension. The unswizzle layout is only defined while the swizzled dim stays in the minormost position, so transposing it is not expressible.","triggerScenarios":"Calling .T or lax.transpose / doing block transposes on a ref that carries a SwizzleTransform, with a permutation where perm[-1] != len(perm) - 1 (last dim is not kept in place).","commonSituations":"Writing a Pallas/Mosaic kernel that transposes a swizzled TMA/WGMMA buffer, or composing ops (reshape+transpose) on tpu swizzled layouts.","solutions":["Restrict transposes to the leading (non-swizzled) dimensions only","Unswizzle (materialize to normal layout) before transposing, then re-swizzle","Redesign the kernel so data is stored pre-transposed instead of transposing in-kernel"],"exampleFix":"// before\nx_t = lax.transpose(x_swizzled, (1, 0))  # perm[-1] != 1 -> ValueError\n// after\nx_n = unswizzle(x_swizzled)\nx_t = lax.transpose(x_n, (1, 0))","handlingStrategy":"validation","validationCode":"def check_transpose_perm(perm):\n    assert perm[-1] == len(perm) - 1, 'Cannot transpose the swizzled (minormost) dim'","typeGuard":null,"tryCatchPattern":"try:\n    y = lax.transpose(x, perm)\nexcept ValueError as e:\n    if 'swizzled dimension' in str(e):\n        x = unswizzle(x); y = lax.transpose(x, perm)\n    else:\n        raise","preventionTips":["Keep the minormost dim fixed when transposing swizzled refs","Unswizzle before arbitrary transposes"],"tags":["jax","pallas","mosaic-gpu","transpose","swizzle"],"backgroundTag":"unsupported-layout-transpose","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}