{"record":{"id":"2864e625d16f3fa0","repo":"jax-ml/jax","slug":"tmem-addr-ref-must-be-a-memref-or-a-pointer-got","errorCode":null,"errorMessage":"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}","messagePattern":"tmem_addr_ref must be a memref or a pointer, got: (.+?)","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/tcgen05.py","lineNumber":1220,"sourceCode":"  @property\n  def packing(self) -> int:\n    return self.layout.vector_length\n\n  def __post_init__(self):\n    self.layout.check_type(self.shape, utils.bitwidth(self.dtype))\n\n  @classmethod\n  def from_alloc(\n      cls,\n      tmem_addr_ref: ir.Value,\n      shape: tuple[int, int],\n      dtype,\n      collective: bool | None = None,\n      layout: TMEMLayout | None = None,\n  ) -> TMEMRef:\n    i32 = ir.IntegerType.get_signless(32)\n    if not isinstance(tmem_addr_ref.type, ir.MemRefType):\n      raise ValueError(f\"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}\")\n    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)\n    if not utils.is_smem_ref(addr_ref_ty):\n      raise ValueError(f\"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}\")\n    if addr_ref_ty.element_type != i32:\n      raise ValueError(f\"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}\")\n    if math.prod(addr_ref_ty.shape) != 1:\n      raise ValueError(f\"tmem_addr_ref must contain a single element, got: {addr_ref_ty}\")\n    i0 = arith.ConstantOp.create_index(0)\n    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)\n    if shape[0] < 32:\n      raise ValueError(f\"TMEM refs must have at least 32 rows, got: {shape[0]}\")\n    if layout is None:\n      if collective is None:\n        raise ValueError(\n            \"collective argument must be provided when TMEM layout is inferred\"\n        )\n      layout = _infer_tmem_layout(shape, collective, packing=1)\n    # TODO: Do we have to do this??","sourceCodeStart":1202,"sourceCodeEnd":1238,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/tcgen05.py#L1202-L1238","documentation":"TMEMRef.from_alloc reads the TMEM base address from a shared-memory memref produced by tmem.alloc. The first check requires tmem_addr_ref to have an MLIR MemRefType (the error text mentions pointer but the code only accepts memrefs); anything else (e.g. an SSA value of another type, a Python int, or an already-loaded scalar) is rejected.","triggerScenarios":"Passing a loaded i32 value, a tensor/ssa value, or a non-memref object as tmem_addr_ref to TMEMRef.from_alloc; passing the result of memref.load instead of the memref itself.","commonSituations":"Hand-writing the alloc plumbing instead of using the helpers; adapting older Mosaic examples where the address handling differed; passing a memref view of wrong type after transformations.","solutions":["Pass the original smem memref holding the allocation address (result of the tmem alloc lowering), not a loaded scalar","If you have a pointer, store it into a 1-element i32 smem memref first and pass that","Check tmem_addr_ref.type prints as memref<...> before calling"],"exampleFix":"# before\naddr = memref.load(alloc_ref, [c0])\nref = tcgen05.TMEMRef.from_alloc(addr, shape, collective=True)\n# after\nref = tcgen05.TMEMRef.from_alloc(alloc_ref, shape, collective=True)  # pass the memref","handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def is_smem_i32_memref(v) -> bool:\n    import ir_module as ir\n    try:\n        return isinstance(v.type, ir.MemRefType)\n    except AttributeError:\n        return False","tryCatchPattern":"try:\n    ref = tcgen05.TMEMRef.from_alloc(addr_ref, shape, collective=c)\nexcept ValueError as e:\n    raise TypeError(f'bad tmem_addr_ref: {e}') from e","preventionTips":["Always pass the memref produced by the tmem alloc helper, never a loaded scalar","Print/assert tmem_addr_ref.type before calling from_alloc when debugging"],"tags":["jax","mosaic","tmem","memref","argument-validation"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}