{"record":{"id":"dd7823eccd185012","repo":"jax-ml/jax","slug":"only-tiledlayout-supports-reductions","errorCode":null,"errorMessage":"Only TiledLayout supports reductions.","messagePattern":"Only TiledLayout supports reductions\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/core.py","lineNumber":1839,"sourceCode":"    object.__setattr__(self, \"kwargs\", frozen_dict.FrozenDict(self.kwargs))\n\n  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:\n    if args or kwargs:\n      raise ValueError(f\"Can't instantiate {self} with arguments.\")\n    return self.layout_cls.to_mgpu(*self.args, **self.kwargs)\n\n\n@dataclasses.dataclass(frozen=True)\nclass ReducedLayout(SomeLayout):\n  layout: SomeLayout\n  axes: Sequence[int]\n\n  def to_mgpu(self, *args, **kwargs) -> mgpu.FragmentedLayout:\n    if args or kwargs:\n      raise ValueError(f\"Can't instantiate {self} with arguments.\")\n    layout = self.layout.to_mgpu()\n    if not isinstance(layout, mgpu.TiledLayout):\n      raise ValueError(\"Only TiledLayout supports reductions.\")\n    return layout.reduce(self.axes)\n\n\nclass Layout(SomeLayout, enum.Enum):\n  #: [m, n] matrix, where m % 64 == 0 == n % 8.\n  WGMMA = enum.auto()\n  WGMMA_8BIT = enum.auto()\n  WGMMA_UPCAST_2X = enum.auto()\n  WGMMA_UPCAST_4X = enum.auto()\n  WGMMA_TRANSPOSED = enum.auto()\n\n  WG_SPLAT = enum.auto()\n  WG_STRIDED = enum.auto()\n\n  TILED = enum.auto()\n\n  TCGEN05 = enum.auto()\n  TCGEN05_TRANSPOSED = enum.auto()","sourceCodeStart":1821,"sourceCodeEnd":1857,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/core.py#L1821-L1857","documentation":"Raised when ReducedLayout.to_mgpu() resolves its inner layout to something other than mgpu.TiledLayout. Reductions (layout.reduce(axes)) are only defined for tiled layouts; reducing an elementwise/fragmented layout such as a transposed WGMMA layout is unsupported.","triggerScenarios":"Constructing ReducedLayout(layout=Layout.WGMMA_TRANSPOSED, axes=[1]) (inner to_mgpu() returns a non-TiledLayout) and then calling to_mgpu().","commonSituations":"Wrapping an arbitrary layout in ReducedLayout when reducing accumulator layouts for collective MMA; refactoring where the inner layout silently changed type.","solutions":["Use a TiledLayout-compatible inner layout (e.g. a TILED/elementwise layout that resolves to mgpu.TiledLayout) before wrapping in ReducedLayout","Perform the reduction manually by reshaping the layout rather than using ReducedLayout"],"exampleFix":"// before\nReducedLayout(Layout.WGMMA_TRANSPOSED, axes=[1]).to_mgpu()\n\n// after\nReducedLayout(Layout.TILED, axes=[1]).to_mgpu()  # inner resolves to TiledLayout","handlingStrategy":"type-guard","validationCode":"inner = reduced.layout.to_mgpu()\nimport jax._src.mosaic_gpu as mgpu\nassert isinstance(inner, mgpu.TiledLayout), 'reductions need a TiledLayout inner'","typeGuard":"def supports_reduction(reduced) -> bool:\n    import jax._src.mosaic_gpu as mgpu\n    return isinstance(reduced.layout.to_mgpu(), mgpu.TiledLayout)","tryCatchPattern":"try:\n    return reduced.to_mgpu()\nexcept ValueError as e:\n    if 'TiledLayout' in str(e):\n        # manual fallback: reshape/tile the layout yourself\n        return manual_reduce(reduced.layout.to_mgpu(), reduced.axes)\n    raise","preventionTips":["Only wrap TiledLayout-producing layouts in ReducedLayout","Check inner layout type before constructing ReducedLayout"],"tags":["jax","pallas","mosaic-gpu","layout","reduction","unsupported-type"],"backgroundTag":"unsupported-layout-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}