{"record":{"id":"f0c3f9b443c502b9","repo":"jax-ml/jax","slug":"sparse-meta-layout-stores-unsupported","errorCode":null,"errorMessage":"Sparse meta layout stores unsupported.","messagePattern":"Sparse meta layout stores unsupported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1479,"sourceCode":"      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)\n    # TODO(apaszke): Support the case where we have a long vector length in the\n    # FA more generally, not just for 2x32b.\n    # TODO(apaszke): Support a wider range of layouts when dealing with unpacking.\n    # 16-bit types are special, because the store instruction can unpack them.\n    elif (\n        value.layout == TMEM_NATIVE_LAYOUT\n        and has_default_layout\n        and ((bitwidth == 16 and packing == 1) or bitwidth == 32)","sourceCodeStart":1461,"sourceCodeEnd":1497,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1461-L1497","documentation":"TensorMem.store refuses to store into TMEM allocated with the sparse metadata layout. The sparse_meta_layout does not faithfully describe the physical TMEM layout written by async_copy_sparse_smem_to_tmem (multiple vector dims would be needed), so a reg->TMEM store could produce inconsistent data.","triggerScenarios":"Allocating TMEM with sparse_meta_layout() for sparse MMA and later calling tmem.store(value) on it.","commonSituations":"Reusing a TMEM allocation for both sparse metadata and general storage; trying to initialize sparse metadata from registers.","solutions":["Do not store into sparse metadata allocations; populate them only via async_copy_sparse_smem_to_tmem","Allocate a separate TMEM buffer for values you need to store","Tag sparse allocations in your kernel code so store is never called on them"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"if tmem.layout == tcgen05.sparse_meta_layout():\n    raise NotImplementedError('cannot store to sparse metadata TMEM')","typeGuard":"def is_sparse_meta(tmem) -> bool:\n    return tmem.layout == tcgen05.sparse_meta_layout()","tryCatchPattern":null,"preventionTips":["Only populate sparse metadata via async_copy_sparse_smem_to_tmem","Keep sparse buffers in dedicated allocations"],"tags":["mosaic","gpu","tcgen05","sparse","not-implemented","tensor-memory","jax"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}