{"record":{"id":"1602ec8d21704352","repo":"jax-ml/jax","slug":"sparse-meta-layout-loads-unsupported","errorCode":null,"errorMessage":"Sparse meta layout loads unsupported.","messagePattern":"Sparse meta layout loads unsupported\\.","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1363,"sourceCode":"      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\n    has_default_layout = self.layout == tmem_default_layout(packing)\n    regs_shape = layout.registers_shape(self.shape)\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 loads unsupported.\")\n    if regs_shape[0] != 1:  # We'll need to issue multiple loads below.\n      raise NotImplementedError(\"Loading multiple row tiles\")\n    if (\n        layout == LAYOUT\n        and self.layout == tmem_default_layout(packing)\n        and is_at_least_16b\n    ):\n      if reduce is not None:\n        raise ValueError(\n            \"Fused load-reduce is not supported for this layout\"\n        )\n      reduced_reg = None\n      registers = _load_32xcols(\n          self.address, columns, self.dtype, packing\n      ).T.reshape(regs_shape)\n    elif layout == self.layout.as_tiled_layout() and packing * bitwidth == 32:\n      # TODO(apaszke): We raise NotImplemented here because technically for some\n      # layouts this does make sense. I think only for those where all","sourceCodeStart":1345,"sourceCodeEnd":1381,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1345-L1381","documentation":"TensorMem.load refuses to load from TMEM carrying the sparse metadata layout. Sparse metadata written by async_copy_sparse_smem_to_tmem has a physical TMEM layout that TiledLayout cannot faithfully describe (it would need multiple vector dims), so loading it back through registers would give inconsistent results.","triggerScenarios":"Allocating TMEM with sparse_meta_layout() (used with async_copy_sparse_smem_to_tmem for sparse MMA) and then calling .load(...) on it.","commonSituations":"Debugging sparse-mma kernels by trying to read sparse metadata back to registers; attempting to inspect/print sparsity metadata via load.","solutions":["Don't load sparse metadata; consume it only via tcgen05.mma with sparse descriptors","If you need to inspect values, copy them through a separate non-sparse TMEM allocation","Track which allocations hold metadata (e.g. a wrapper class) so load is never called on them"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"if tmem.layout == tcgen05.sparse_meta_layout():\n    raise NotImplementedError('cannot load sparse metadata; consume via tcgen05.mma')","typeGuard":"def is_sparse_meta(tmem) -> bool:\n    return tmem.layout == tcgen05.sparse_meta_layout()","tryCatchPattern":null,"preventionTips":["Wrap sparse-meta allocations in a small class that disables load","Keep metadata buffers separate from value buffers in kernel structure"],"tags":["mosaic","gpu","tcgen05","sparse","tensor-memory","not-implemented","jax"],"backgroundTag":"unsupported-layout-operation","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}