{"record":{"id":"d9371043fd254c0f","repo":"jax-ml/jax","slug":"expected-a-memref-type-but-got-ref","errorCode":null,"errorMessage":"Expected a memref type but got {ref}","messagePattern":"Expected a memref type but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/utils.py","lineNumber":2252,"sourceCode":"def tmem() -> ir.Attribute:\n  \"\"\"Returns the attribute for the TMEM memory space.\"\"\"\n  return ir.Attribute.parse(\"#mosaic_gpu.tmem\")\n\n\ndef smem_cluster() -> ir.Attribute:\n  \"\"\"Returns the attribute for the cluster SMEM memory space.\"\"\"\n  return ir.Attribute.parse(\"#mosaic_gpu.smem_cluster\")\n\n\ndef is_smem_ref(ref: ir.Value | ir.Type) -> bool:\n  \"\"\"Returns true if the input mem ref or memref type points to SMEM.\n\n  If the input is not at all of a memref type, raises a ValueError.\n  \"\"\"\n  if isinstance(ref, ir.Value):\n    ref = ref.type\n  if not isinstance(ref, ir.MemRefType):\n    raise ValueError(f\"Expected a memref type but got {ref}\")\n  ref = ir.MemRefType(ref)\n  return ref.memory_space is not None and ref.memory_space == smem()\n\n\ndef is_tmem_ref(ref: ir.Value | ir.Type) -> bool:\n  \"\"\"Returns true if the input mem ref or memref type points to TMEM.\n\n  If the input is not at all of a memref type, raises a ValueError.\n  \"\"\"\n  if isinstance(ref, ir.Value):\n    ref = ref.type\n  if not isinstance(ref, ir.MemRefType):\n    raise ValueError(f\"Expected a memref type but got {ref}\")\n  ref = ir.MemRefType(ref)\n  return ref.memory_space is not None and ref.memory_space == tmem()\n\n\ndef is_cluster_smem_ref(ref: ir.Value | ir.Type) -> bool:","sourceCodeStart":2234,"sourceCodeEnd":2270,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/utils.py#L2234-L2270","documentation":"is_smem_ref (and its docstring) accept an ir.Value or ir.Type and check whether it is a MemRefType located in shared memory (memory_space == smem()). Anything that is not a memref — tensor, vector, function result types — raises this ValueError as a misuse of the predicate.","triggerScenarios":"Calling is_smem_ref with a tensor-typed value, a vector value, or any non-memref ir.Type, e.g. checking the result of an op that produces a tensor.","commonSituations":"Feeding values from tensor-level (pre-lowering) IR into this memref-level helper; passing a module-level type or a function type accidentally; refactors where a value's type changed from memref to tensor.","solutions":["Ensure the value has been lowered to buffers (memref) before calling is_smem_ref","Check isinstance(value.type, ir.MemRefType) first and handle the non-memref case explicitly","Fix the producer op to emit a memref result (e.g. use bufferization/to_memref)"],"exampleFix":"# before\nin_smem = is_smem_ref(tensor_val)\n# after\nin_smem = isinstance(tensor_val.type, ir.MemRefType) and is_smem_ref(tensor_val)","handlingStrategy":"type-guard","validationCode":"if isinstance(ref, ir.Value):\n    ref = ref.type\nif isinstance(ref, ir.MemRefType):\n    ok = is_smem_ref(ref)\nelse:\n    ok = False  # or handle non-memref case","typeGuard":"def is_memref(v_or_t) -> bool:\n    t = v_or_t.type if isinstance(v_or_t, ir.Value) else v_or_t\n    return isinstance(t, ir.MemRefType)","tryCatchPattern":null,"preventionTips":["Lower tensors to memrefs before passing to memref-level helpers","Wrap predicates with isinstance checks on ir.MemRefType"],"tags":["mosaic-gpu","memref","type-validation","shared-memory"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}