jax-ml/jax · error · NotImplementedError

Unsupported layout: {src.layout}

Error message

Unsupported layout: {src.layout}

What it means

The multi-dim reduction lowering requires the source value to be in a TiledLayout; fragment (register) layouts are not supported as the reduction source.

Source

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

) -> 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
    allocation_size = ir.IntegerAttr(op.attributes["scratch_size"]).value * 8 // utils.bitwidth(dtype)
    scratch = _slice_smem(
        ir.MemRefType.get([allocation_size], dtype, memory_space=utils.smem()),
        ir.IntegerAttr(op.attributes["offset"]).value,
        ctx.smem_requested_bytes,
    )
  else:
    scratch = None

  match op_kind:
    case vector.CombiningKind.ADD:
      result = src.reduce("add", op.reduction_dims[0], scratch)
      result += acc
    case vector.CombiningKind.MAXSI | vector.CombiningKind.MAXUI | vector.CombiningKind.MAXIMUMF:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. layout_cast the source to a TiledLayout before the reduction
2. Restructure so the reduction input comes from a load with tiled layout

Example fix

// before
r = multi_dim_reduction(acc_fragment, acc)
// after
tiled = layout_cast(acc_fragment, tiled_layout)
r = multi_dim_reduction(tiled, acc)

assert isinstance(src.layout, fa.TiledLayout), 'cast source to TiledLayout before reduction'

def is_tiled(fa_value) -> bool:
    return isinstance(fa_value.layout, fa.TiledLayout)

When it happens

Trigger: Passing a value with a fragment layout (e.g. straight out of a matmul) to multi_dim_reduction.

Common situations: Reducing a wgmma accumulator without first casting to a tiled layout.

Related errors

• Output layout {out_layout} must match the accumulator 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/132842625a5b56c8. Report an issue: GitHub.