{"record":{"id":"da55b6bbf01113df","repo":"jax-ml/jax","slug":"memref-loadop-does-not-support-transforms-op","errorCode":null,"errorMessage":"memref.LoadOp does not support transforms: {op}","messagePattern":"memref\\.LoadOp does not support transforms: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2363,"sourceCode":"      unwrap_transformed_memref(op.src, in_transforms_attr),\n      new_reassociation,\n  )\n  return [wrap_transformed_memref(result, op.result.type, out_transforms_attr)]\n\n\n@_register_lowering(memref.LoadOp)\ndef _memref_load_op_lowering_rule(\n    ctx: LoweringContext, op: memref.LoadOp\n) -> Sequence[ir.Value]:\n  \"\"\"Lowering rule for memref.LoadOp.\n\n  Loads are never transformed so this rule is mostly just a pass-through.\n  \"\"\"\n  del ctx\n\n  in_transforms = inference_utils.in_transforms(op)[0]\n  if in_transforms:\n    raise NotImplementedError(f\"memref.LoadOp does not support transforms: {op}\")\n\n  new_load_op = memref.LoadOp(\n      memref=unwrap_transformed_memref(op.memref, in_transforms),\n      indices=op.indices,\n      nontemporal=op.nontemporal,\n  )\n  return [new_load_op.result]\n\n\n@_register_lowering(memref.StoreOp, support_warp_semantics=True)\ndef _memref_store_op_lowering_rule(\n    ctx: LoweringContext, op: memref.StoreOp\n) -> Sequence[ir.Value]:\n  \"\"\"Lowering rule for memref.StoreOp.\n\n  Stores are never transformed so this rule is mostly just a pass-through.\n  \"\"\"\n  del ctx","sourceCodeStart":2345,"sourceCodeEnd":2381,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2345-L2381","documentation":"memref.load is a scalar access and Mosaic's lowering is a pass-through: it explicitly rejects any non-empty transform annotation on the loaded memref, since transformed layouts can't be indexed directly.","triggerScenarios":"Emitting memref.load where inference_utils.in_transforms(op)[0] is non-empty, i.e. loading from a memref that still carries tiling/swizzle transforms.","commonSituations":"Loading individual elements from a tiled/swizzled smem tensor without unwrapping the transform first; using low-level load instead of Mosaic's tensor loads on transformed buffers.","solutions":["Unwrap/discard the transform before loading (load from the base memref via unwrap_transformed_memref)","Use load_tensor/store_tensor ops which do support transformed memrefs","Avoid tiling memrefs you intend to access element-wise"],"exampleFix":"// before\nv = t.memref.load(tiled_ref, idx)\n// after\nbase = unwrap_transformed_memref(tiled_ref, transforms)\nv = t.memref.load(base, idx)","handlingStrategy":"validation","validationCode":"assert not inference_utils.in_transforms(op)[0], 'load does not support transforms'","typeGuard":"def loadable_untransformed(op) -> bool:\n    return not inference_utils.in_transforms(op)[0]","tryCatchPattern":null,"preventionTips":["Unwrap transforms before scalar loads","Use load_tensor for transformed buffers"],"tags":["jax","mosaic-gpu","memref-load","transforms","not-implemented"],"backgroundTag":"unsupported-layout-transform","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}