{"record":{"id":"132842625a5b56c8","repo":"jax-ml/jax","slug":"unsupported-layout-src-layout","errorCode":null,"errorMessage":"Unsupported layout: {src.layout}","messagePattern":"Unsupported layout: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":897,"sourceCode":") -> Sequence[ir.Value]:\n  [in_layout, acc_layout] = inference_utils.in_layouts(op)\n  [out_layout] = inference_utils.out_layouts(op)\n  if out_layout != acc_layout:\n    raise ValueError(\n        f\"Output layout {out_layout} must match the accumulator layout\"\n        f\" {acc_layout}\"\n    )\n\n  if len(op.reduction_dims) != 1:\n    raise NotImplementedError(\"Only 1 reduction dimension is supported.\")\n\n  op_kind = _combining_kind(op.kind)\n  is_signed = _is_reduction_signed(op_kind)\n  src = _fragmented_array_from_ir(op.source, in_layout, is_signed)\n  acc = _fragmented_array_from_ir(op.acc, acc_layout, is_signed)\n\n  if not isinstance(src.layout, fa.TiledLayout):\n    raise NotImplementedError(f\"Unsupported layout: {src.layout}\")\n  reduced_dim = src.layout.tiling.tile_dimension(op.reduction_dims[0])\n  if any(reduced_dim[d] for d in src.layout.partitioned_warp_dims):\n    # cross-warp reductions require scratch space.\n    dtype = op.source.type.element_type\n    allocation_size = ir.IntegerAttr(op.attributes[\"scratch_size\"]).value * 8 // utils.bitwidth(dtype)\n    scratch = _slice_smem(\n        ir.MemRefType.get([allocation_size], dtype, memory_space=utils.smem()),\n        ir.IntegerAttr(op.attributes[\"offset\"]).value,\n        ctx.smem_requested_bytes,\n    )\n  else:\n    scratch = None\n\n  match op_kind:\n    case vector.CombiningKind.ADD:\n      result = src.reduce(\"add\", op.reduction_dims[0], scratch)\n      result += acc\n    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:","sourceCodeStart":879,"sourceCodeEnd":915,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L879-L915","documentation":"The multi-dim reduction lowering requires the source value to be in a TiledLayout; fragment (register) layouts are not supported as the reduction source.","triggerScenarios":"Passing a value with a fragment layout (e.g. straight out of a matmul) to multi_dim_reduction.","commonSituations":"Reducing a wgmma accumulator without first casting to a tiled layout.","solutions":["layout_cast the source to a TiledLayout before the reduction","Restructure so the reduction input comes from a load with tiled layout"],"exampleFix":"// before\nr = multi_dim_reduction(acc_fragment, acc)\n// after\ntiled = layout_cast(acc_fragment, tiled_layout)\nr = multi_dim_reduction(tiled, acc)","handlingStrategy":"type-guard","validationCode":"assert isinstance(src.layout, fa.TiledLayout), 'cast source to TiledLayout before reduction'","typeGuard":"def is_tiled(fa_value) -> bool:\n    return isinstance(fa_value.layout, fa.TiledLayout)","tryCatchPattern":null,"preventionTips":["layout_cast matmul accumulators before reducing"],"tags":["jax","mosaic","gpu","reduction","layout"],"backgroundTag":"unsupported-layout","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}