{"record":{"id":"c4e02c711123f7c0","repo":"jax-ml/jax","slug":"memref-storeop-does-not-support-transforms-op","errorCode":null,"errorMessage":"memref.StoreOp does not support transforms: {op}","messagePattern":"memref\\.StoreOp does not support transforms: (.+?)","errorType":"exception","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2385,"sourceCode":"      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\n\n  in_transforms = inference_utils.in_transforms(op)[0]\n  if in_transforms:\n    raise NotImplementedError(f\"memref.StoreOp does not support transforms: {op}\")\n\n  memref.StoreOp(\n      value=op.value,\n      memref=unwrap_transformed_memref(op.memref, in_transforms),\n      indices=op.indices,\n      nontemporal=op.nontemporal,\n  )\n  return []\n\n\n@_register_lowering(mgpu.TmemAllocOp)\ndef _tmem_alloc_op_lowering_rule(\n    ctx: LoweringContext, op: mgpu.TmemAllocOp\n) -> Sequence[ir.Value]:\n  \"\"\"Lowering rule for mgpu.TmemAllocOp.\"\"\"\n  ctx.check_collective(op)\n\n  output_shape = ir.MemRefType(op.result.type).shape","sourceCodeStart":2367,"sourceCodeEnd":2403,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2367-L2403","documentation":"memref.store, like load, is a scalar pass-through in Mosaic lowering and rejects any non-empty transform annotation on the target memref.","triggerScenarios":"Emitting memref.store where the stored-to memref has non-empty in_transforms (tiled/swizzled layout still attached).","commonSituations":"Writing scalar results into a tiled shared-memory buffer directly; mixing element-wise stores with tensor-mem transforms in a fused kernel.","solutions":["Unwrap the transform and store to the underlying base memref","Use store_tensor for transformed buffers","Keep memrefs used for scalar stores untransformed"],"exampleFix":"// before\nt.memref.store(val, tiled_ref, idx)\n// after\nbase = unwrap_transformed_memref(tiled_ref, transforms)\nt.memref.store(val, base, idx)","handlingStrategy":"validation","validationCode":"assert not inference_utils.in_transforms(op)[0], 'store does not support transforms'","typeGuard":"def storable_untransformed(op) -> bool:\n    return not inference_utils.in_transforms(op)[0]","tryCatchPattern":null,"preventionTips":["Unwrap transforms before scalar stores","Use store_tensor for transformed buffers"],"tags":["jax","mosaic-gpu","memref-store","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"}