{"record":{"id":"533037ca8e3fcf6d","repo":"jax-ml/jax","slug":"loads-from-tmem-layout-self-layout-to-register-l","errorCode":null,"errorMessage":"Loads from TMEM layout {self.layout} to register layout {layout} are not supported","messagePattern":"Loads from TMEM layout (.+?) to register layout (.+?) are not supported","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1434,"sourceCode":"      )\n      assert raw_registers.shape[0] == 4\n      registers = np.concatenate([raw_registers[:2], raw_registers[2:]], axis=1)\n      registers = registers.T.reshape(regs_shape)\n    elif (\n        layout == fa_m64_collective_layout(columns)\n        and self.layout == tmem_m64_collective_layout(columns, packing)\n        and is_at_least_16b\n    ):\n      if reduce is not None:\n        raise ValueError(\"Fused load-reduce is not supported for this layout\")\n      reduced_reg = None\n      regs_shape = layout.registers_shape(self.shape)\n      # We take half the columns, because they are split over halves of TMEM.\n      registers = _load_32xcols(\n          self.address, columns // 2, self.dtype, packing\n      ).reshape(regs_shape)\n    else:\n      raise ValueError(\n          f\"Loads from TMEM layout {self.layout} to register layout\"\n          f\" {layout} are not supported\"\n      )\n    result = fa.FragmentedArray(\n        _registers=registers, _layout=layout, _is_signed=is_signed\n    )\n    if reduce is None:\n      # The None assignments in the branches let us use the linter to ensure\n      # that we didn't forget to handle reduce in any of the cases.\n      assert reduced_reg is None\n      return result\n    reduced_layout = layout.reduce((len(layout.base_tile_shape) - 1,))\n    assert reduced_layout.vector_length == 1\n    reduced_regs_shape = reduced_layout.registers_shape(self.shape[:-1])\n    assert math.prod(reduced_regs_shape) == 1\n    reduced_result = fa.FragmentedArray(\n        _registers=np.asarray(reduced_reg, dtype=object).reshape(reduced_regs_shape),\n        _layout=reduced_layout,","sourceCodeStart":1416,"sourceCodeEnd":1452,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1416-L1452","documentation":"TensorMem.load fell through all supported layout cases (default layout, native tiled, wgmma half-lane, m64 collective) and none matched the (self.layout, layout) pair. This is the catch-all for unsupported TMEM-to-register layout conversions.","triggerScenarios":"Calling tmem.load(layout) with a register layout that doesn't match any supported pair for the TMEM's layout — e.g. a custom TiledLayout, a transposed layout, or an unsupported packing/bitwidth combination (packing*bitwidth != 32 on the native path).","commonSituations":"Custom register layouts for new kernels; 8-bit or 64-bit element types where packing constraints fail; layout mismatch after kernel refactors.","solutions":["Read the branches above the raise to pick a supported layout pair (LAYOUT + default TMEM layout, as_tiled_layout with packing*bitwidth==32, WGMMA_LAYOUT + half-lane, m64 collective)","Ensure packing satisfies 32-bit column constraints (e.g. packing = 32 // bitwidth)","Construct TMEM with a standard layout via tmem_default_layout / tmem_half_lane_layout / tmem_m64_collective_layout helpers"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"def load_pair_supported(tmem, layout, bitwidth):\n    dl = tcgen05.tmem_default_layout(tmem.packing)\n    return ((layout == tcgen05.LAYOUT and tmem.layout == dl and bitwidth in (16,32))\n        or (layout == tmem.layout.as_tiled_layout() and tmem.packing*bitwidth == 32)\n        or (layout == fa.WGMMA_LAYOUT and tmem.layout == tcgen05.tmem_half_lane_layout(tmem.shape[1], tmem.packing))\n        or (layout == tcgen05.fa_m64_collective_layout(tmem.shape[1]) and tmem.layout == tcgen05.tmem_m64_collective_layout(tmem.shape[1], tmem.packing)))","typeGuard":null,"tryCatchPattern":"try:\n    arr, red = tmem.load(layout, reduce=reduce)\nexcept ValueError as e:\n    if 'not supported' in str(e):\n        layout = choose_supported_layout(tmem); arr, _ = tmem.load(layout)\n    else: raise","preventionTips":["Allocate TMEM with layouts produced by the tcgen05 helper functions","Keep a table of supported (tmem_layout, register_layout) pairs in your kernel library"],"tags":["mosaic","gpu","tcgen05","layout","tensor-memory","jax"],"backgroundTag":"unsupported-layout-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}