jax-ml/jax · error · ValueError

{ref} has a layout {layout_attr} that does not match the exp

Error message

{ref} has a layout {layout_attr} that does not match the expected layout {expected_layout}.

What it means

The conversion_cast attached to a TMEM operand carries a layout attribute that must exactly equal the expected TiledLayout computed by layout inference. Mismatch means the layout recorded when the value was produced differs from what the consuming tcgen05 op requires.

Source

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

  """
  if not isinstance(ref.type, ir.MemRefType):
    raise ValueError(f"{ref} is not a memref.")
  mem_ref_ty = ir.MemRefType(ref.type)

  if mem_ref_ty.memory_space != utils.tmem():
    raise ValueError(
        f"{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM."
    )

  i32 = ir.IntegerType.get_signless(32)
  conversion_cast, [tmem_addr] = _undo_conversion_cast(ref, [i32])

  assert mem_ref_ty.rank == 2
  shape = cast(tuple[int, int], tuple(mem_ref_ty.shape))
  el_ty = mem_ref_ty.element_type
  layout_attr = conversion_cast.attributes["layout"]
  if layout_attr != expected_layout:
    raise ValueError(
        f"{ref} has a layout {layout_attr} that does not match the expected"
        f" layout {expected_layout}."
    )
  layout = layouts_lib.from_layout_attr(layout_attr)
  assert isinstance(layout, fa.TiledLayout)
  tmem_layout = tcgen05.TMEMLayout(
      layout.tiling, layout.warp_dims, layout.lane_dims, layout.vector_dim
  )
  return tcgen05.TMEMRef(tmem_addr, shape, el_ty, tmem_layout)


def _tmem_ref_to_ir(ref: tcgen05.TMEMRef, ty: ir.MemRefType) -> ir.Value:
  """Returns an IR value from a TMEMRef."""
  conversion_cast = builtin.UnrealizedConversionCastOp([ty], [ref.address])
  conversion_cast.attributes["layout"] = layouts_lib.to_layout_attr(ref.layout)
  return conversion_cast.result

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Print both layouts (layout_attr vs expected_layout) and align your cast target with the expected tmem_layout of the consumer
2. Fix the producer (slice_tmem / cast) so its resulting layout attribute matches; don't manually fabricate layout attrs
3. Update to a Mosaic version where the layouts of paired ops are kept consistent

Example fix

// before
ref = mgpu.tmem_layout_cast(acc, target_layout=some_other_layout)
out = mgpu.tcgen05_mma(a, b, ref)
// after
ref = mgpu.tmem_layout_cast(acc, target_layout=tmem_layout(acc))  # match consumer expectation
out = mgpu.tcgen05_mma(a, b, ref)

expected = tmem_layout(ref)  # what consumer expects
actual = conversion_cast_of(ref).attributes['layout']
assert actual == expected

try:
    out = mgpu.tcgen05_mma(a, b, acc)
except ValueError as e:
    if 'does not match the expected layout' in str(e):
        # realign cast target with tmem_layout and retry
        acc = mgpu.tmem_layout_cast(acc, tmem_layout(acc))
        out = mgpu.tcgen05_mma(a, b, acc)
    else:
        raise

When it happens

Trigger: Using a layout-cast or slice op that changes the tiling of a TMEM value, then feeding it to a tcgen05 op whose expected layout (from tmem_layout() of the operand) differs; mixing tmem_layout_cast with incompatible source/target layouts.

Common situations: Hand-composed layout pipelines where tmem_layout_cast targets a non-batching layout the MMA op doesn't expect; version changes to default tcgen05 accumulator layouts.

Related errors

• {ref} is not a memref.
• {ref} has a memory space {mem_ref_ty.memory_space} that is n
• Input layouts {in_layouts} do not match yield layouts {yield
• Output layouts {out_layouts} do not match condition layouts
• Rule cannot handle an op with vector operands or results: {o

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