{"record":{"id":"7c0b6c4a8e70fe41","repo":"jax-ml/jax","slug":"cannot-apply-tiling-transform-str-with-swizzle","errorCode":null,"errorMessage":"Cannot apply {tiling_transform_str} with {swizzle_str} to memref with shape {source.shape}.","messagePattern":"Cannot apply (.+?) with (.+?) to memref with shape (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2165,"sourceCode":"  source = ValueSite(op, VariableType.OPERAND, 0)\n  dest = ValueSite(op, VariableType.RESULT, 0)\n  var = ctx.producer_ref(source)\n  smem_transforms = _extract_smem_transforms_from_custom_transform_attrs(\n      op.transforms\n  )\n\n  if not cs.is_valid_assignment(var, smem_transforms):\n    tiling_transform_str = (\n        f\"tiling {smem_transforms.tiling}\"\n        if smem_transforms.tiling\n        else \"empty tiling\"\n    )\n    swizzle_str = (\n        f\"{smem_transforms.swizzle} swizzle\"\n        if smem_transforms.swizzle\n        else \"no swizzle\"\n    )\n    raise ValueError(\n        f\"Cannot apply {tiling_transform_str} with {swizzle_str} to memref with\"\n        f\" shape {source.shape}.\"\n    )\n  assignments: dict[cs.Variable, cs.Constant] = {var: smem_transforms}\n  return cs.ConstraintSystem(assignments=assignments), {var: [source, dest]}\n\n\ndef _vector_value_sites_and_constraints_for_async_ops(\n    op: mgpu.AsyncLoadOp | mgpu.AsyncStoreOp | mgpu.AsyncPrefetchOp,\n) -> tuple[ValueSitesForVariable, list[cs.Constraint]]:\n  values_sites: ValueSitesForVariable = dict()\n  constraints: list[cs.Constraint] = []\n\n  match op:\n    case mgpu.AsyncLoadOp():\n      base_operand_index = 3\n    case mgpu.AsyncStoreOp():\n      base_operand_index = 2","sourceCodeStart":2147,"sourceCodeEnd":2183,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2147-L2183","documentation":"When applying with_transforms, Mosaic checks that the tiling vector (and optional swizzle) evenly divides the source memref's shape. If any dimension of the memref is not a multiple of the corresponding tiling (or is smaller than the tile), the transforms cannot be applied and a ValueError is raised describing the mismatch.","triggerScenarios":"Calling with_transforms (or a custom primitive carrying smem transforms) where e.g. tiling [128, 128] with swizzle 128 is applied to a memref of shape [64, 64], or a dimension not divisible by the tile/swizzle factor.","commonSituations":"Hard-coded tile/swizzle sizes (common 128B swizzle) combined with small or odd-shaped tail tiles; kernels parameterized by problem size where small N triggers the check; changing swizzle mode without resizing smem buffers.","solutions":["Make the smem/tile shape a multiple of the tiling vector and swizzle factor (e.g. pad the last tile dimension up to the swizzle bytes)","Derive the tiling vector from the memref shape: tile dims should divide the corresponding memref dims","Reduce the swizzle (e.g. from 128B to 64B/32B) or disable swizzle for small tiles"],"exampleFix":"# before: ValueError — 64 not divisible by 128\nt = mgpu.smem_transforms(tiling=[128, 128], swizzle=128)\nout = mgpu.with_transforms(smem_shaped_64x64, t)\n\n# after\nt = mgpu.smem_transforms(tiling=[64, 64], swizzle=64)\nout = mgpu.with_transforms(smem_shaped_64x64, t)","handlingStrategy":"validation","validationCode":"tile, sw = smem_transforms.tile, smem_transforms.swizzle or 1\nassert all(d % t == 0 and d % sw == 0 for d, t in zip(shape, tile)), 'pad memref to multiple of tiling/swizzle'","typeGuard":"def transforms_fit(shape, tiling, swizzle=None) -> bool:\n    sw = swizzle or 1\n    return all(d % t == 0 and d % sw == 0 for d, t in zip(shape, tiling))","tryCatchPattern":null,"preventionTips":["Pad smem tile dimensions up to the swizzle byte width (32/64/128)","Compute tiling vectors from the actual buffer shape, not global constants"],"tags":["jax","mosaic-gpu","smem-transforms","swizzle","tiling","shape-mismatch"],"backgroundTag":"shape-dimension-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}