{"record":{"id":"d94d5a80a1dc3c71","repo":"jax-ml/jax","slug":"tmem-layout-self-layout-is-not-supported","errorCode":null,"errorMessage":"TMEM layout {self.layout} is not supported","messagePattern":"TMEM layout (.+?) is not supported","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1335,"sourceCode":"      layout: fa.TiledLayout | None = None,\n      is_signed: bool | None = None,\n      reduce: LoadReduceOp | None = None,  # Reduction operator for the minor dimension.\n  ) -> fa.FragmentedArray | tuple[fa.FragmentedArray, fa.FragmentedArray]:\n    packing = self.packing\n    bitwidth = utils.bitwidth(self.dtype)\n    is_at_least_16b = bitwidth in {16, 32}\n    columns = self.shape[1]\n    if layout is None:\n      if is_at_least_16b and self.layout == tmem_default_layout(packing):\n        layout = LAYOUT\n      elif is_at_least_16b and packing <= columns // 2 and self.layout == tmem_half_lane_layout(columns, packing):\n        layout = fa.WGMMA_LAYOUT\n      elif is_at_least_16b and columns % 16 == 0 and self.layout == tmem_m64_collective_layout(columns, packing):\n        layout = fa_m64_collective_layout(columns)\n      elif packing * bitwidth == 32:\n        layout = self.layout.as_tiled_layout()\n      else:\n        raise ValueError(f\"TMEM layout {self.layout} is not supported\")\n    if reduce is not None:\n      if isinstance(self.dtype, ir.IntegerType) and bitwidth == 32:\n        if reduce not in (\"min\", \"max\"):\n          raise ValueError(\n              \"Unsupported reduction for i32. Only min and max are supported,\"\n              f\" got: {reduce}\"\n          )\n        if not is_signed:\n          reduce = \"abs\" + reduce  # type: ignore\n      elif isinstance(self.dtype, ir.F32Type):\n        if reduce not in (\"min\", \"max\", \"absmin\", \"absmax\"):\n          raise ValueError(\n              \"Unsupported reduction for f32. Only min, max, absmin, and\"\n              f\" absmax are supported, got: {reduce}\"\n          )\n      else:\n        raise ValueError(f\"Unsupported dtype for reduction: {self.dtype}\")\n","sourceCodeStart":1317,"sourceCodeEnd":1353,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1317-L1353","documentation":"tcgen05.load converts TMEM to registers using a fragment-array layout matched to the TMEM layout and element bitwidth. Only specific combos are supported: WGMMA_LAYOUT, the m64 collective pair, or the layout as-is when packing*bitwidth==32. Any other TMEM layout raises 'TMEM layout … is not supported'.","triggerScenarios":"Calling tcgen05.load on a TMEM ref with a custom/relayouted TMEMLayout that doesn't match the supported patterns, or where packing * element_bitwidth != 32 (e.g. packing=1 with f32, or packing=2 with i8).","commonSituations":"Building custom TMEMLayouts instead of the provided constructors; using packing values inconsistent with dtype width (e.g. packing=4 with 16-bit types gives 64 != 32).","solutions":["Use the canonical layouts: tmem_default_layout, tmem_half_lane_layout, or tmem_m64_collective_layout","Choose packing so packing * bitwidth == 32 (f32->1, f16/bf16->2, 8-bit->4, 4-bit->8) or one of the other supported branches","If you need an unsupported layout, relayout in registers after a supported load"],"exampleFix":"# before\nref = tcgen05.TMEMRef(..., layout=custom_layout)  # packing*bitwidth != 32\nout = tcgen05.load(ref)\n# after\nref = tcgen05.TMEMRef(..., layout=tcgen05.tmem_default_layout(packing=2))  # bf16: 2*16==32\nout = tcgen05.load(ref)","handlingStrategy":"validation","validationCode":"bitwidth = ref.dtype.width\npacking = ref.layout.vector_size  # or however packing is tracked\nassert packing * bitwidth == 32 or ref.layout in (known_supported,)","typeGuard":"def load_supported(ref) -> bool:\n    bitwidth = getattr(ref.dtype, 'width', 32)\n    return ref.layout in SUPPORTED_TMEM_LAYOUTS_FOR(ref) or ref.layout_packing * bitwidth == 32","tryCatchPattern":"try:\n    out = tcgen05.load(ref)\nexcept ValueError as e:\n    if 'not supported' in str(e):\n        ref = ref.to_layout(tcgen05.tmem_default_layout(packing=32 // bitwidth))\n        out = tcgen05.load(ref)\n    else:\n        raise","preventionTips":["Always create TMEM refs with the three canonical layout constructors","Set packing = 32 // element_bitwidth so the as_tiled_layout fallback branch applies"],"tags":["jax","mosaic","tmem","load","layout","unsupported-operation"],"backgroundTag":"unsupported-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}