{"record":{"id":"166ccb9543474b21","repo":"jax-ml/jax","slug":"output-layout-out-layout-must-match-the-accumula","errorCode":null,"errorMessage":"Output layout {out_layout} must match the accumulator layout {acc_layout}","messagePattern":"Output layout (.+?) must match the accumulator layout (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":883,"sourceCode":"      result = a.reduce(\"add\", axes, scratch)\n    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:\n      result = a.reduce(\"max\", axes, scratch)\n    case vector.CombiningKind.MINUI | vector.CombiningKind.MINSI | vector.CombiningKind.MINIMUMF:\n      result = a.reduce(\"min\", axes, scratch)\n    case _:\n      raise NotImplementedError(f\"Unsupported reduction kind: {op.kind}\")\n  assert isinstance(result.layout, fa.WGSplatFragLayout)\n  return [result.registers.item()]\n\n\n@_register_lowering(vector.MultiDimReductionOp)\ndef _vector_multi_dim_reduction_op_lowering_rule(\n    ctx: LoweringContext, op: vector.MultiDimReductionOp\n) -> 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","sourceCodeStart":865,"sourceCodeEnd":901,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L865-L901","documentation":"For vector.multi_dim_reduction the output layout must equal the accumulator's input layout; Mosaic cannot produce a result whose layout differs from the accumulator's.","triggerScenarios":"Emitting multi_dim_reduction where the op's out_layout attr differs from the acc operand's in_layout attr.","commonSituations":"Layout inference assigns a different layout to the result (e.g. after a layout_cast on the accumulator) causing mismatch at lowering time.","solutions":["Insert a layout_cast on the result (or accumulator) so both share the same layout","Annotate the reduction op's result layout to match the accumulator"],"exampleFix":"// before\nr = multi_dim_reduction(src, acc)\n// after\nr = multi_dim_reduction(src, acc)\nr = layout_cast(r, acc_layout)","handlingStrategy":"validation","validationCode":"assert out_layout == acc_layout, 'annotate result layout to match accumulator'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep accumulator and reduction result layouts identical; cast after"],"tags":["jax","mosaic","gpu","reduction","layout"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}