{"record":{"id":"167bd6a1d9830052","repo":"jax-ml/jax","slug":"storing-from-register-layout-value-layout-to-tme","errorCode":null,"errorMessage":"Storing from register layout {value.layout} to TMEM layout {self.layout} is not supported","messagePattern":"Storing from register layout (.+?) to TMEM layout (.+?) is not supported","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1517,"sourceCode":"      _store_32xcols_native(self.address, value.registers.reshape(-1), packing)\n    elif (\n        value.layout == fa.WGMMA_LAYOUT\n        and self.layout == tmem_half_lane_layout(self.shape[1], packing=packing)\n        and is_at_least_16b\n    ):\n      registers = value.registers.T.reshape(2, -1)\n      registers = np.concatenate(np.split(registers, 2, axis=1), axis=0)\n      _store_32xcols(self.address, registers, packing)\n    elif (\n        value.layout == fa_m64_collective_layout(self.shape[1])\n        and self.layout == tmem_m64_collective_layout(\n            self.shape[1], packing=packing\n        )\n        and is_at_least_16b\n    ):\n      _store_32xcols(self.address, value.registers.reshape(4, -1), packing)\n    else:\n      raise ValueError(\n          f\"Storing from register layout {value.layout} to TMEM layout\"\n          f\" {self.layout} is not supported\"\n      )\n\n  def _debug_print(self) -> None:\n    i32 = ir.IntegerType.get_signless(32)\n    num_cols = self.layout.cols_in_shape(self.shape, utils.bitwidth(self.dtype))\n    lane = arith.remui(utils.thread_idx(), arith.constant(i32, utils.WARPGROUP_SIZE))\n    for c in range(num_cols):\n      ptr = _tmem_addr_to_ptr(arith.addi(self.address, arith.constant(i32, c)))\n      i32_vec = ir.VectorType.get((1,), i32)\n      vec_val = nvvm.tcgen05_ld(i32_vec, nvvm.Tcgen05LdStShape.SHAPE_32X32B, ptr)\n      val = vector.extract(vec_val, [], [0])\n      dtype_bitwidth = utils.bitwidth(self.dtype)\n      full_packing = 32 // dtype_bitwidth\n      if self.packing == 1:\n        if dtype_bitwidth < 32:\n          val = arith.trunci(ir.IntegerType.get_signless(dtype_bitwidth), val)","sourceCodeStart":1499,"sourceCodeEnd":1535,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1499-L1535","documentation":"Catch-all in TensorMem.store: the (value.layout, self.layout) pair matched none of the supported store paths (default layout, native tiled, wgmma half-lane, m64 collective). The register layout cannot be lowered into this TMEM layout.","triggerScenarios":"Storing a FragmentedArray whose TiledLayout doesn't pair with the TMEM allocation's layout, or with unsupported packing/bitwidth (e.g. 8-bit elements where is_at_least_16b fails).","commonSituations":"Custom tiled layouts; sub-16-bit element types; layout drift after refactoring a kernel's TMEM allocation strategy.","solutions":["Relayout the value to match one of the supported pairs (e.g. tmem_default_layout with LAYOUT registers)","Use a standard TMEM layout helper for allocation","Ensure element bitwidth is 16 or 32 and packing satisfies the 32-bit column constraint"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert utils.bitwidth(tmem.dtype) in (16, 32), 'store needs >=16b elements'\nassert value.layout == tmem.layout.as_tiled_layout() or tmem.layout == tcgen05.tmem_default_layout(packing=tmem.packing), 'unsupported store pair'","typeGuard":null,"tryCatchPattern":"try:\n    tmem.store(value)\nexcept ValueError:\n    value = value.relayout(tmem.layout.as_tiled_layout()); tmem.store(value)","preventionTips":["Use tcgen05 layout helpers for TMEM allocation","Relayout values to match the TMEM layout before storing"],"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"}