{"record":{"id":"f580eb94e6b696b4","repo":"jax-ml/jax","slug":"no-swizzle-is-not-supported-f580eb","errorCode":null,"errorMessage":"No swizzle is not supported","messagePattern":"No swizzle is not supported","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/wgmma.py","lineNumber":335,"sourceCode":"    a: fa.FragmentedArray | ir.Value,\n    b: ir.Value,\n    *,\n    swizzle: int = 128,\n):\n  \"\"\"Perform acc += a @ b using the WGMMA instruction.\n\n  `a` may be passed in registers, or as a memref. `b` must be a memref.\n\n  The expected (logical) memref shapes are:\n    a: (m // tile_m, k // tile_k, tile_m, tile_k)\n    b: (k // tile_k, n // tile_n, tile_k, tile_n).\n\n  While the shapes may be physically transposed, when considering the row-major\n  physical shape, the tile dimensions must be the two minor dimensions and must\n  have the shape (8, S) where S = swizzle // bytewidth(element_type).\n  \"\"\"\n  if swizzle == 16:\n    raise NotImplementedError(\"No swizzle is not supported\")\n  # Step 1. Establish the shape and element type of the operation.\n  if not isinstance(b.type, ir.MemRefType):\n    raise ValueError(f\"B must be a memref, got: {b.type}\")\n  bf16 = ir.BF16Type.get()\n  f32 = ir.F32Type.get()\n  f16 = ir.F16Type.get()\n  i32 = ir.IntegerType.get_signless(32)\n  i8 = ir.IntegerType.get_signless(8)\n  f8e5m2 = ir.Float8E5M2Type.get()\n  f8e4m3fn = ir.Float8E4M3FNType.get()\n  (k, n), element_type = mma_utils.tiled_memref_shape(b)\n  if a_in_regs := isinstance(a, fa.FragmentedArray):\n    m, k2 = a.shape\n    element_type2 = a.mlir_dtype\n    if element_type2 not in {f16, bf16, i8, f8e5m2, f8e4m3fn}:\n      raise ValueError(\n          \"Only f16, bf16, i8, f8e5m2, f8e4m3fn are supported for A \"\n          f\"in registers, got {element_type2}\"","sourceCodeStart":317,"sourceCodeEnd":353,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/wgmma.py#L317-L353","documentation":"At wgmma.py:335, the wgmma wrapper rejects swizzle == 16 (i.e. 'no swizzle'): the shared-memory descriptor path in Mosaic only supports the 32/64/128-byte swizzle modes.","triggerScenarios":"Calling wgmma.wgmma(..., swizzle=16); hit from lowering rules or the FlashAttention-style kernels (compute_qk, compute_pv, etc.) when the SMEM layout was allocated with 16-byte swizzle.","commonSituations":"Passing utils.swizzle(16) as the swizzle mode when allocating B; porting TMA descriptors configured for no swizzle; assuming all power-of-two swizzles are allowed.","solutions":["Use utils.swizzle(32/64/128) for the B operand's SMEM allocation","Re-allocate/re-layout SMEM with a supported swizzle before the wgmma call","Check utils.swizzle constants in your JAX version and pick from the supported set"],"exampleFix":"# before\nlayout = utils.swizzle(16)\nb = utils.memref_alloc(..., layout)\nacc = wgmma.wgmma(a, b, acc, swizzle=16)\n# after\nlayout = utils.swizzle(32)\nacc = wgmma.wgmma(a, b, acc, swizzle=32)","handlingStrategy":"validation","validationCode":"assert swizzle in (32, 64, 128), f'swizzle 16 unsupported, got {swizzle}'","typeGuard":null,"tryCatchPattern":"try:\n    wgmma.wgmma(a, b, acc, swizzle=swizzle)\nexcept NotImplementedError:\n    wgmma.wgmma(a, b, acc, swizzle=32)","preventionTips":["Use only utils.swizzle(32/64/128) modes","Match TMA descriptor swizzle to one of the supported modes"],"tags":["jax","mosaic-gpu","wgmma","swizzle","smem"],"backgroundTag":"unsupported-config-value","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}