{"record":{"id":"72705c57a8277ae3","repo":"jax-ml/jax","slug":"ref-is-not-a-memref","errorCode":null,"errorMessage":"{ref} is not a memref.","messagePattern":"(.+?) is not a memref\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2459,"sourceCode":"\n  output_shape = ir.MemRefType(op.tmem_ref.type).shape\n  ncols = output_shape[1] // packing\n\n  with utils.when(ctx.single_warp_per_block_predicate):\n    tcgen05.tmem_dealloc(tmem_addr, ncols, collective, exact=False)\n\n  return []\n\n\ndef _tmem_ref_from_ir(\n    ref: ir.Value, expected_layout: ir.Attribute\n) -> tcgen05.TMEMRef:\n  \"\"\"Returns a TMEMRef from an IR value.\n\n  Throws an error if the annotated layout does not match the expected layout.\n  \"\"\"\n  if not isinstance(ref.type, ir.MemRefType):\n    raise ValueError(f\"{ref} is not a memref.\")\n  mem_ref_ty = ir.MemRefType(ref.type)\n\n  if mem_ref_ty.memory_space != utils.tmem():\n    raise ValueError(\n        f\"{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM.\"\n    )\n\n  i32 = ir.IntegerType.get_signless(32)\n  conversion_cast, [tmem_addr] = _undo_conversion_cast(ref, [i32])\n\n  assert mem_ref_ty.rank == 2\n  shape = cast(tuple[int, int], tuple(mem_ref_ty.shape))\n  el_ty = mem_ref_ty.element_type\n  layout_attr = conversion_cast.attributes[\"layout\"]\n  if layout_attr != expected_layout:\n    raise ValueError(\n        f\"{ref} has a layout {layout_attr} that does not match the expected\"\n        f\" layout {expected_layout}.\"","sourceCodeStart":2441,"sourceCodeEnd":2477,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2441-L2477","documentation":"Raised by Mosaic GPU's TMEM lowering when a value passed where a tensor-memory (TMEM) memref is expected does not have an ir.MemRefType. The _tmem_ref_from_ir helper validates that the IR value is a memref before converting it into a tcgen05.TMEMRef for Blackwell tcgen05 ops.","triggerScenarios":"Calling ops that lower to tcgen05 TMEM ops (tcgen05_mma, async_copy to TMEM, tmem_layout_cast, slice_tmem, print_layout) with an operand whose MLIR type is not a MemRefType — e.g. a tensor value or a scalar that was never allocated into TMEM.","commonSituations":"Building Mosaic kernels that pass a register-resident tensor or an incorrectly-produced value into a tmem argument slot; wiring custom primitives or layouts where the operand producer was skipped.","solutions":["Check the producer of the offending operand — it must be an mgpu.mem_ref / TMEM allocation producing a memref type","Ensure you wrapped the value with the proper TMEM allocation op (e.g. local_tensor with tmem memory space) before passing it","Verify no layout-inference step dropped the memref typing on the operand chain"],"exampleFix":"// before\nacc = my_tensor  # tensor, not TMEM memref\nmgpu.tcgen05_mma(a, b, acc)\n// after\nacc = mgpu.local_alloc(my_tensor, memory_space='tmem')\nmgpu.tcgen05_mma(a, b, acc)","handlingStrategy":"type-guard","validationCode":"from jaxlib.mlir import ir\ndef is_memref(v):\n    return isinstance(v.type, ir.MemRefType)","typeGuard":"def is_tmem_operand(v) -> bool:\n    from jaxlib.mlir import ir\n    t = getattr(v, 'type', None)\n    return isinstance(t, ir.MemRefType)","tryCatchPattern":null,"preventionTips":["Always allocate TMEM operands via mgpu local-tensor APIs with tmem memory space","Never feed raw tensors into tcgen05 operand slots"],"tags":["mosaic","gpu","tmem","mlir","type-validation"],"backgroundTag":"invalid-operand-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}