{"record":{"id":"b12edaf1a46de805","repo":"jax-ml/jax","slug":"no-swizzle-is-not-supported","errorCode":null,"errorMessage":"No swizzle is not supported","messagePattern":"No swizzle is not supported","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":193,"sourceCode":"      raise NotImplementedError(f\"Unsupported input dtype: {ty}\")\n  return _create_scaled_instr_descriptor(get_input_encoding, *args, **kwargs)\n\n\ndef mma(\n    d: TMEMRef,\n    a: ir.Value | TMEMRef,\n    b: ir.Value,\n    *,\n    a_swizzle: int = 128,\n    b_swizzle: int = 128,\n    a_scale: TMEMRef | None = None,\n    b_scale: TMEMRef | None = None,\n    a_sparse_metadata: TMEMRef | None = None,\n    accumulate: ir.Value | bool = True,\n    collective: bool = False,\n) -> None:\n  if a_swizzle == 16 or b_swizzle == 16:\n    raise NotImplementedError(\"No swizzle is not supported\")\n  i8 = ir.IntegerType.get_signless(8)\n  i32 = ir.IntegerType.get_signless(32)\n  if isinstance(accumulate, bool):\n    accumulate = arith.constant(ir.IntegerType.get_signless(1), accumulate)\n  num_cta = 2 if collective else 1\n  if (is_scaled := a_scale is not None) != (b_scale is not None):\n    raise ValueError(\"Either none or both scales should be provided\")\n  is_sparse = a_sparse_metadata is not None\n  if is_scaled and is_sparse:\n    if isinstance(a, TMEMRef):\n      raise NotImplementedError(\n          \"A in TMEM unsupported for block-scaled sparse matmuls\"\n      )\n\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  (k, n), b_element_type = mma_utils.tiled_memref_shape(b)","sourceCodeStart":175,"sourceCodeEnd":211,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L175-L211","documentation":"tcgen05 mma() rejects swizzle mode 16 (the 'no swizzle'/32-byte-atomic layout) because that shared-memory layout is not supported by the tcgen05 tensor-core path. Only 32B/64B/128B swizzles work.","triggerScenarios":"Calling mma(a_swizzle=16) or mma(b_swizzle=16).","commonSituations":"Porting Hopper WGMMA kernels that used the 16-byte (none) swizzle mode, or defaulting swizzle parameters to 16 from older code.","solutions":["Use a_swizzle/b_swizzle of 32, 64, or 128","Re-layout shared memory operands with a supported swizzle mode before the mma"],"exampleFix":"# before\nmma(acc, a, b, a_swizzle=16, b_swizzle=16)\n\n# after\nmma(acc, a, b, a_swizzle=128, b_swizzle=128)","handlingStrategy":"validation","validationCode":"assert a_swizzle != 16 and b_swizzle != 16, 'swizzle 16 unsupported; use 32/64/128'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Never reuse swizzle=16 (none) configs from WGMMA code","Default shared-memory layouts to 128B swizzle"],"tags":["tcgen05","matmul","swizzle","shared-memory"],"backgroundTag":"unsupported-swizzle-mode","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}