{"record":{"id":"aa00750529ac0c36","repo":"jax-ml/jax","slug":"swizzle-self-swizzle-requires-the-trailing-dimen","errorCode":null,"errorMessage":"Swizzle {self.swizzle} requires the trailing dimension to be of size {swizzle_elems}, but got shape: {x.shape}","messagePattern":"Swizzle (.+?) requires the trailing dimension to be of size (.+?), but got shape: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1185,"sourceCode":"@dataclasses.dataclass(frozen=True)\nclass SwizzleTransform(state_types.Transform):\n  swizzle: int\n\n  def __post_init__(self):\n    if self.swizzle not in {32, 64, 128}:\n      raise ValueError(\n          f\"Swizzle {self.swizzle} is not supported. Only 32, 64 and 128 are\"\n          \" accepted.\"\n      )\n\n  def transform_type(\n      self, x: jax_core.AbstractValue\n  ) -> jax_core.AbstractValue:\n    match x:\n      case jax_core.ShapedArray():\n        swizzle_elems = (self.swizzle * 8) // dtypes.itemsize_bits(x.dtype)\n        if swizzle_elems != x.shape[-1]:\n          raise ValueError(\n              f\"Swizzle {self.swizzle} requires the trailing dimension to be of\"\n              f\" size {swizzle_elems}, but got shape: {x.shape}\"\n          )\n        return x\n      case state_types.AbstractRef():\n        return x.update(inner_aval=self.transform_type(x.inner_aval))\n      case _:\n        raise NotImplementedError(f\"Unsupported type: {x}\")\n\n  def undo(self, x: jax_core.AbstractValue) -> state_types.Transform:\n    return UnswizzleRef(self.swizzle)\n\n\n@tree_util.register_dataclass\n@dataclasses.dataclass(frozen=True)\nclass UnswizzleRef(state_types.Transform):\n  swizzle: int = jax.tree.static()\n","sourceCodeStart":1167,"sourceCodeEnd":1203,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1167-L1203","documentation":"SwizzleTransform.transform_type checks that the trailing (minormost) dimension size equals the swizzle width in elements (swizzle_bytes*8 / itemsize_bits). If the last dim doesn't match, the swizzled layout would be ill-defined and a ValueError with the expected size is raised.","triggerScenarios":"Applying SwizzleTransform(128) to a block whose last dimension isn't swizzle_elems long — e.g. fp32 (32-bit) block with trailing dim != 32 for swizzle 128, or padding changing the last dim.","commonSituations":"Swizzled TMA layouts where block shapes/dtypes don't match the swizzle: bf16 with swizzle 64 needs trailing dim 32; using padded or ragged trailing dimensions.","solutions":["Pad/reshape the block so the trailing dim equals (swizzle*8)//itemsize_bits","Choose the swizzle that matches your block: for trailing dim N and dtype, pick 32/64/128 so N == swizzle_elems","Avoid applying the swizzle transform to non-matching intermediate shapes; apply it only at the storage layout level"],"exampleFix":"# before\n# bf16 block shape (..., 16) with SwizzleTransform(128) -> needs 64 elems\nref = ... SwizzleTransform(128)\n\n# after\n# pad trailing dim to 64, or use SwizzleTransform(32) for 16 bf16 elems (16*16=256b -> use 32B swizzle)\nref = ... SwizzleTransform(32)  # matches trailing dim 16 for bf16? verify: 32*8/16=16","handlingStrategy":"validation","validationCode":"import jax.numpy as jnp\nfrom jax._src import dtypes\nexpected = (swizzle * 8) // dtypes.itemsize_bits(dtype)\nassert x.shape[-1] == expected, f'trailing dim must be {expected}'","typeGuard":"def swizzle_matches(swizzle, shape, dtype): return shape[-1] == (swizzle * 8) // dtypes.itemsize_bits(dtype)","tryCatchPattern":null,"preventionTips":["Match swizzle bytes to block trailing dim and dtype","Pad blocks to swizzle-aligned trailing dims"],"tags":["jax","pallas","mosaic-gpu","swizzle","layout","shape-mismatch"],"backgroundTag":"layout-shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}