{"record":{"id":"d974e76afae25ee1","repo":"jax-ml/jax","slug":"memref-cast-tmem-layouts-must-be-identical-for-bot","errorCode":null,"errorMessage":"memref.cast tmem layouts must be identical for both input and output but got {in_tmem_layout=} and {out_tmem_layout=}","messagePattern":"memref\\.cast tmem layouts must be identical for both input and output but got (.+?) and (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2102,"sourceCode":"  if memory_space == utils.smem():\n    [in_transforms] = inference_utils.in_transforms(op)\n    [out_transforms] = inference_utils.out_transforms(op)\n    if in_transforms != out_transforms:\n      raise NotImplementedError(\n          \"memref.cast transforms must have identical transforms for both \"\n          f\"input and output but got {in_transforms=} and {out_transforms=}\"\n      )\n    result = memref.cast(\n        transform_type(ir.MemRefType(op.result.type), out_transforms),\n        unwrap_transformed_memref(op.source, in_transforms),\n    )\n    return [wrap_transformed_memref(result, op.result.type, out_transforms)]\n\n  if memory_space == utils.tmem():\n    [in_tmem_layout] = inference_utils.in_tmem_layouts(op)\n    [out_tmem_layout] = inference_utils.out_tmem_layouts(op)\n    if in_tmem_layout != out_tmem_layout:\n      raise NotImplementedError(\n          \"memref.cast tmem layouts must be identical for both input and\"\n          f\" output but got {in_tmem_layout=} and {out_tmem_layout=}\"\n      )\n    return [_tmem_ref_to_ir(_tmem_ref_from_ir(op.source, in_tmem_layout),\n                            op.result.type)]\n\n  raise NotImplementedError(\n      f\"Unsupported memory space when lowering memref.cast: {memory_space}\"\n  )\n\n\ndef _permutation_to_affine_map_attr(\n    permutation: Sequence[int],\n) -> ir.AffineMapAttr:\n  return ir.AffineMapAttr.get(ir.AffineMap.get_permutation(permutation))\n\n\n@_register_lowering(memref.TransposeOp, support_warp_semantics=True)","sourceCodeStart":2084,"sourceCodeEnd":2120,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2084-L2120","documentation":"For tensor-memory (tmem) casts in Mosaic GPU, the input and output TMEM layout annotations must match exactly; the lowering just re-wraps the same tmem reference and cannot reconcile different layouts.","triggerScenarios":"memref.cast with memory_space == utils.tmem() where the op's in_tmem_layouts and out_tmem_layouts annotations differ (different tmem packing/dp word layout).","commonSituations":"Working with tcgen05 tensor memory on Blackwell: copying a tmem ref and simultaneously trying to reinterpret its layout; mismatch usually comes from layout inference defaults when a cast result type was constructed by hand.","solutions":["Use the same TMEM layout annotation on both operands (pass the source layout through)","Allocate the destination with the matching layout instead of relayouting via cast","Print in/out tmem layouts via inference_utils.in_tmem_layouts(op) to find the divergence"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"[in_l] = inference_utils.in_tmem_layouts(op)\n[out_l] = inference_utils.out_tmem_layouts(op)\nassert in_l == out_l, 'tmem layouts must match for cast'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Reuse the source tmem layout on cast results","Let inference fill tmem layouts"],"tags":["jax","mosaic-gpu","tmem","cast","layout"],"backgroundTag":"tmem-layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}