{"record":{"id":"65992ddfb06477d6","repo":"jax-ml/jax","slug":"loading-multiple-row-tiles","errorCode":null,"errorMessage":"Loading multiple row tiles","messagePattern":"Loading multiple row tiles","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1365,"sourceCode":"          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\n      # dimensions that map to columns map only to columns.\n      if reduce is not None:","sourceCodeStart":1347,"sourceCodeEnd":1383,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1347-L1383","documentation":"TensorMem.load only supports tensors whose register shape has a single row tile (regs_shape[0] == 1). When the requested register layout needs multiple row tiles, the load would require issuing several loads, which is not implemented.","triggerScenarios":"Calling tmem.load(layout) where layout.registers_shape(tmem.shape)[0] != 1 — e.g. tall/narrow register layouts spanning more than 128 lanes worth of rows for the TMEM allocation.","commonSituations":"Choosing a FragmentedArray register layout with more row tiles than the TMEM shape supports; mismatched layouts when refactoring kernels from one shape to another.","solutions":["Use a register layout whose registers_shape(tmem.shape)[0] == 1 (e.g. LAYOUT or wgmma-compatible layouts)","Check layout.registers_shape(tmem.shape) before calling load","Resize the TMEM allocation or restructure the kernel so one row tile suffices"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"regs_shape = layout.registers_shape(tmem.shape)\nassert regs_shape[0] == 1, f'load supports one row tile, got {regs_shape[0]}'","typeGuard":"def load_supported(tmem, layout) -> bool:\n    return layout.registers_shape(tmem.shape)[0] == 1","tryCatchPattern":null,"preventionTips":["Validate layout.registers_shape against TMEM shape in kernel setup","Prefer standard layouts (LAYOUT, WGMMA_LAYOUT) for loads"],"tags":["mosaic","gpu","tcgen05","layout","tensor-memory","not-implemented","jax"],"backgroundTag":"unsupported-layout-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}