{"record":{"id":"7149dc9b7c00d7b4","repo":"jax-ml/jax","slug":"memref-cast-transforms-must-have-identical-transfo","errorCode":null,"errorMessage":"memref.cast transforms must have identical transforms for both input and output but got {in_transforms=} and {out_transforms=}","messagePattern":"memref\\.cast transforms must have identical transforms 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":2088,"sourceCode":"  unoffseted_out_ty = ir.MemRefType.get(\n      out_ty.shape,\n      out_ty.element_type,\n      memory_space=out_ty.memory_space,\n      layout=ir.StridedLayoutAttr.get(0, out_strides),\n  )\n\n  if unoffseted_in_ty != unoffseted_out_ty:\n    raise NotImplementedError(\n        \"Only support memref.cast where the input and output types are the \"\n        f\"same up to offset, but got {in_ty=} and {out_ty=}.\"\n    )\n\n  memory_space = ir.MemRefType(op.result.type).memory_space\n  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),","sourceCodeStart":2070,"sourceCodeEnd":2106,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2070-L2106","documentation":"When casting a shared-memory memref, Mosaic requires the transform annotations on input and output to be exactly identical, because smem casts are pure pointer reinterpretations. Any divergence in tiling/swizzle annotations is rejected.","triggerScenarios":"memref.cast into utils.smem() where inference_utils.in_transforms(op) != out_transforms(op), e.g. annotating the result with a different tile shape or swizzle than the source.","commonSituations":"Manually setting transforms on a cast result while the source carries inferred transforms; or re-tiling a tensor in smem via cast instead of an explicit transform op.","solutions":["Copy the source's transform annotation verbatim to the cast result","Use a dedicated transform op (tile/swizzle) to change layout rather than cast","Inspect both transform attrs with inference_utils.in/out_transforms(op) and diff them"],"exampleFix":"// before\nout = t.memref.cast(in_smem, ty_with_different_transforms)\n// after\n[in_t] = in_transforms  # reuse\nout = t.memref.cast(in_smem, transform_type(ty, in_t))","handlingStrategy":"validation","validationCode":"[in_t] = inference_utils.in_transforms(op)\n[out_t] = inference_utils.out_transforms(op)\nassert in_t == out_t, 'smem cast transforms must match'","typeGuard":"def smem_cast_transforms_match(op) -> bool:\n    return inference_utils.in_transforms(op) == inference_utils.out_transforms(op)","tryCatchPattern":null,"preventionTips":["Copy source transform annotation to cast results","Avoid changing layout via cast"],"tags":["jax","mosaic-gpu","memref","cast","shared-memory","transforms"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}