{"record":{"id":"5e09fbab8dc60234","repo":"jax-ml/jax","slug":"ref-has-a-layout-layout-attr-that-does-not-mat","errorCode":null,"errorMessage":"{ref} has a layout {layout_attr} that does not match the expected layout {expected_layout}.","messagePattern":"(.+?) has a layout (.+?) that does not match the expected layout (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2475,"sourceCode":"  \"\"\"\n  if not isinstance(ref.type, ir.MemRefType):\n    raise ValueError(f\"{ref} is not a memref.\")\n  mem_ref_ty = ir.MemRefType(ref.type)\n\n  if mem_ref_ty.memory_space != utils.tmem():\n    raise ValueError(\n        f\"{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM.\"\n    )\n\n  i32 = ir.IntegerType.get_signless(32)\n  conversion_cast, [tmem_addr] = _undo_conversion_cast(ref, [i32])\n\n  assert mem_ref_ty.rank == 2\n  shape = cast(tuple[int, int], tuple(mem_ref_ty.shape))\n  el_ty = mem_ref_ty.element_type\n  layout_attr = conversion_cast.attributes[\"layout\"]\n  if layout_attr != expected_layout:\n    raise ValueError(\n        f\"{ref} has a layout {layout_attr} that does not match the expected\"\n        f\" layout {expected_layout}.\"\n    )\n  layout = layouts_lib.from_layout_attr(layout_attr)\n  assert isinstance(layout, fa.TiledLayout)\n  tmem_layout = tcgen05.TMEMLayout(\n      layout.tiling, layout.warp_dims, layout.lane_dims, layout.vector_dim\n  )\n  return tcgen05.TMEMRef(tmem_addr, shape, el_ty, tmem_layout)\n\n\ndef _tmem_ref_to_ir(ref: tcgen05.TMEMRef, ty: ir.MemRefType) -> ir.Value:\n  \"\"\"Returns an IR value from a TMEMRef.\"\"\"\n  conversion_cast = builtin.UnrealizedConversionCastOp([ty], [ref.address])\n  conversion_cast.attributes[\"layout\"] = layouts_lib.to_layout_attr(ref.layout)\n  return conversion_cast.result\n\n","sourceCodeStart":2457,"sourceCodeEnd":2493,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2457-L2493","documentation":"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.","triggerScenarios":"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.","commonSituations":"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.","solutions":["Print both layouts (layout_attr vs expected_layout) and align your cast target with the expected tmem_layout of the consumer","Fix the producer (slice_tmem / cast) so its resulting layout attribute matches; don't manually fabricate layout attrs","Update to a Mosaic version where the layouts of paired ops are kept consistent"],"exampleFix":"// before\nref = mgpu.tmem_layout_cast(acc, target_layout=some_other_layout)\nout = mgpu.tcgen05_mma(a, b, ref)\n// after\nref = mgpu.tmem_layout_cast(acc, target_layout=tmem_layout(acc))  # match consumer expectation\nout = mgpu.tcgen05_mma(a, b, ref)","handlingStrategy":"validation","validationCode":"expected = tmem_layout(ref)  # what consumer expects\nactual = conversion_cast_of(ref).attributes['layout']\nassert actual == expected","typeGuard":null,"tryCatchPattern":"try:\n    out = mgpu.tcgen05_mma(a, b, acc)\nexcept ValueError as e:\n    if 'does not match the expected layout' in str(e):\n        # realign cast target with tmem_layout and retry\n        acc = mgpu.tmem_layout_cast(acc, tmem_layout(acc))\n        out = mgpu.tcgen05_mma(a, b, acc)\n    else:\n        raise","preventionTips":["Derive cast targets from tmem_layout() of the consumer operand, never hardcode","Add debug prints of layout attrs before tcgen05 ops during kernel bring-up"],"tags":["mosaic","gpu","tmem","layout-mismatch"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}