{"record":{"id":"b1f00ec73532d0e7","repo":"jax-ml/jax","slug":"fused-load-reduce-is-not-supported-for-this-layout","errorCode":null,"errorMessage":"Fused load-reduce is not supported for this layout","messagePattern":"Fused load-reduce is not supported for this layout","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1372,"sourceCode":"    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:\n        raise NotImplementedError(\n            \"Fused load-reduce is not supported for this layout\"\n        )\n      assert len(layout.base_tile_shape) == 2\n      cols = math.prod(regs_shape) * packing\n      flat_registers, reduced_reg = _load_32xcols_native(\n          self.address, cols, self.dtype, packing, packing, reduce=None","sourceCodeStart":1354,"sourceCodeEnd":1390,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1354-L1390","documentation":"When loading from a default TMEM layout (>=16-bit elements) directly into the LAYOUT register layout via _load_32xcols, no fused reduction is performed, so passing a reduce argument is rejected.","triggerScenarios":"tmem.load(LAYOUT, reduce=...) where self.layout == tmem_default_layout(packing) and element bitwidth is 16 or 32.","commonSituations":"Writing generic load-reduce helper code that always passes reduce; switching from a layout that supports fused reduce (e.g. native tiled) to the default TMEM layout.","solutions":["Call load without reduce and reduce in registers afterwards","Use a TMEM/register layout pair that supports fused load-reduce (the _load_32xcols_native path)","Only pass reduce when the layout conditions permit it"],"exampleFix":"// before\narr = tmem.load(tcgen05.LAYOUT, reduce='max')\n// after\narr = tmem.load(tcgen05.LAYOUT)\nred = arr.reduce('max', axis=0)","handlingStrategy":"validation","validationCode":"if reduce is not None and tmem.layout == tcgen05.tmem_default_layout(tmem.packing):\n    reduce = None  # fused reduce unsupported on this path","typeGuard":null,"tryCatchPattern":"try:\n    arr, red = tmem.load(layout, reduce=reduce)\nexcept ValueError:\n    arr, _ = tmem.load(layout); red = manual_reduce(arr)","preventionTips":["Only pass reduce for layout pairs documented to support it","Write a thin wrapper: load_maybe_reduce(tmem, layout, reduce) that checks the pair"],"tags":["mosaic","gpu","tcgen05","reduction","layout","jax"],"backgroundTag":"unsupported-layout-combination","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}