jax-ml/jax · error · ValueError

Output layout {out_layout} must match the accumulator layout

Error message

Output layout {out_layout} must match the accumulator layout {acc_layout}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:883

      result = a.reduce("add", axes, scratch)
    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:
      result = a.reduce("max", axes, scratch)
    case vector.CombiningKind.MINUI | vector.CombiningKind.MINSI | vector.CombiningKind.MINIMUMF:
      result = a.reduce("min", axes, scratch)
    case _:
      raise NotImplementedError(f"Unsupported reduction kind: {op.kind}")
  assert isinstance(result.layout, fa.WGSplatFragLayout)
  return [result.registers.item()]


@_register_lowering(vector.MultiDimReductionOp)
def _vector_multi_dim_reduction_op_lowering_rule(
    ctx: LoweringContext, op: vector.MultiDimReductionOp
) -> Sequence[ir.Value]:
  [in_layout, acc_layout] = inference_utils.in_layouts(op)
  [out_layout] = inference_utils.out_layouts(op)
  if out_layout != acc_layout:
    raise ValueError(
        f"Output layout {out_layout} must match the accumulator layout"
        f" {acc_layout}"
    )

  if len(op.reduction_dims) != 1:
    raise NotImplementedError("Only 1 reduction dimension is supported.")

  op_kind = _combining_kind(op.kind)
  is_signed = _is_reduction_signed(op_kind)
  src = _fragmented_array_from_ir(op.source, in_layout, is_signed)
  acc = _fragmented_array_from_ir(op.acc, acc_layout, is_signed)

  if not isinstance(src.layout, fa.TiledLayout):
    raise NotImplementedError(f"Unsupported layout: {src.layout}")
  reduced_dim = src.layout.tiling.tile_dimension(op.reduction_dims[0])
  if any(reduced_dim[d] for d in src.layout.partitioned_warp_dims):
    # cross-warp reductions require scratch space.
    dtype = op.source.type.element_type

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Insert a layout_cast on the result (or accumulator) so both share the same layout
2. Annotate the reduction op's result layout to match the accumulator

Example fix

// before
r = multi_dim_reduction(src, acc)
// after
r = multi_dim_reduction(src, acc)
r = layout_cast(r, acc_layout)

assert out_layout == acc_layout, 'annotate result layout to match accumulator'

When it happens

Trigger: Emitting multi_dim_reduction where the op's out_layout attr differs from the acc operand's in_layout attr.

Common situations: Layout inference assigns a different layout to the result (e.g. after a layout_cast on the accumulator) causing mismatch at lowering time.

Related errors

• Unsupported layout: {src.layout}
• Fused load-reduce is not supported for this layout
• {op} has an unsupported layout: {out_layout_attr}
• Expected TiledLayout, got {type(layout)}
• Unsupported reduction kind: {op.kind}

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/166ccb9543474b21. Report an issue: GitHub.