{"record":{"id":"e4b24eee86ec6a39","repo":"jax-ml/jax","slug":"stored-array-has-layout-value-layout-but-tmem-s","errorCode":null,"errorMessage":"Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout","messagePattern":"Stored array has layout (.+?), but TMEM stores expect a TiledLayout","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1471,"sourceCode":"        _is_signed=is_signed,\n    )\n    return result, reduced_result\n\n  def store(self, value: fa.FragmentedArray):\n    if not isinstance(value, fa.FragmentedArray):\n      raise TypeError(f\"TMEM stores expect a FragmentedArray, got: {value}\")\n    if value.shape != self.shape:\n      raise ValueError(\n          f\"Stored array has shape {value.shape}, but TMEM has shape\"\n          f\" {self.shape}\"\n      )\n    if value.mlir_dtype != self.dtype:\n      raise ValueError(\n          f\"Stored array has dtype {value.mlir_dtype}, but TMEM has dtype\"\n          f\" {self.dtype}\"\n      )\n    if not isinstance(value.layout, fa.TiledLayout):\n      raise TypeError(f\"Stored array has layout {value.layout}, but TMEM stores expect a TiledLayout\")\n    # TODO(olechwierowicz): `sparse_meta_layout()` does not really describe the\n    # actual TMEM layout of the result of `async_copy_sparse_smem_to_tmem`.\n    # As a result storing through SMEM -> Reg -> TMEM is not equivalent to\n    # SMEM -> TMEM. We raise in this case to prevent inconsistent behaviour.\n    # This restriction can be lifted if `TiledLayout` supports multiple\n    # vector dims.\n    if self.layout == sparse_meta_layout():\n      raise NotImplementedError(\"Sparse meta layout stores unsupported.\")\n    packing = self.packing\n    has_default_layout = self.layout == tmem_default_layout(packing=packing)\n    bitwidth = utils.bitwidth(self.dtype)\n    is_at_least_16b = bitwidth in {16, 32}\n    if value.layout == LAYOUT and has_default_layout and is_at_least_16b:\n      _store_32xcols(\n          self.address, value.registers.T.reshape((4, -1)), packing\n      )\n    elif value.layout == self.layout.as_tiled_layout() and packing * bitwidth == 32:\n      _store_32xcols_native(self.address, value.registers.reshape(-1), packing)","sourceCodeStart":1453,"sourceCodeEnd":1489,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1453-L1489","documentation":"TensorMem.store requires the incoming FragmentedArray's layout to be a TiledLayout (register layouts like LAYOUT or WGMMA_LAYOUT are not accepted). Store lowering only knows how to move tiled-layout registers into TMEM.","triggerScenarios":"tmem.store(value) where value.layout is a strided register layout such as tcgen05.LAYOUT or fa.WGMMA_LAYOUT instead of a TiledLayout instance.","commonSituations":"Storing a freshly computed mma result (which has a register layout) back to TMEM; re-layouting via relayout to a tiled layout before store is required.","solutions":["Relayout the array first: value = value.relayout(some_tiled_layout) so value.layout is a fa.TiledLayout","Store via SMEM (round-trip through shared memory) if a suitable tiled layout is unavailable","Check isinstance(value.layout, fa.TiledLayout) before storing"],"exampleFix":"// before\ntmem.store(mma_result)  # layout is WGMMA_LAYOUT\n// after\ntiled = mma_result.relayout(tmem.layout.as_tiled_layout())\ntmem.store(tiled)","handlingStrategy":"type-guard","validationCode":"if not isinstance(value.layout, fa.TiledLayout):\n    value = value.relayout(tmem.layout.as_tiled_layout())","typeGuard":"def has_tiled_layout(v) -> bool:\n    return isinstance(v.layout, fa.TiledLayout)","tryCatchPattern":"try:\n    tmem.store(value)\nexcept TypeError:\n    tmem.store(value.relayout(tmem.layout.as_tiled_layout()))","preventionTips":["Relayout mma results to a tiled layout before any TMEM store"],"tags":["mosaic","gpu","tcgen05","layout","type-check","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"}