jax-ml/jax · error · ValueError

tmem_addr_ref must be a memref or a pointer, got: {tmem_addr

Error message

tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/tcgen05.py:1220

  @property
  def packing(self) -> int:
    return self.layout.vector_length

  def __post_init__(self):
    self.layout.check_type(self.shape, utils.bitwidth(self.dtype))

  @classmethod
  def from_alloc(
      cls,
      tmem_addr_ref: ir.Value,
      shape: tuple[int, int],
      dtype,
      collective: bool | None = None,
      layout: TMEMLayout | None = None,
  ) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    if not isinstance(tmem_addr_ref.type, ir.MemRefType):
      raise ValueError(f"tmem_addr_ref must be a memref or a pointer, got: {tmem_addr_ref.type}")
    addr_ref_ty = ir.MemRefType(tmem_addr_ref.type)
    if not utils.is_smem_ref(addr_ref_ty):
      raise ValueError(f"tmem_addr_ref must be in shared memory, got: {addr_ref_ty}")
    if addr_ref_ty.element_type != i32:
      raise ValueError(f"tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}")
    if math.prod(addr_ref_ty.shape) != 1:
      raise ValueError(f"tmem_addr_ref must contain a single element, got: {addr_ref_ty}")
    i0 = arith.ConstantOp.create_index(0)
    tmem_addr = memref.load(tmem_addr_ref, [i0] * addr_ref_ty.rank)
    if shape[0] < 32:
      raise ValueError(f"TMEM refs must have at least 32 rows, got: {shape[0]}")
    if layout is None:
      if collective is None:
        raise ValueError(
            "collective argument must be provided when TMEM layout is inferred"
        )
      layout = _infer_tmem_layout(shape, collective, packing=1)
    # TODO: Do we have to do this??

View on GitHub (pinned to 1e1c6a8fc0)

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

Example fix

# before
addr = memref.load(alloc_ref, [c0])
ref = tcgen05.TMEMRef.from_alloc(addr, shape, collective=True)
# after
ref = tcgen05.TMEMRef.from_alloc(alloc_ref, shape, collective=True)  # pass the memref

Defensive patterns

Strategy: type-guard

Type guard

def is_smem_i32_memref(v) -> bool:
    import ir_module as ir
    try:
        return isinstance(v.type, ir.MemRefType)
    except AttributeError:
        return False

Try / catch

try:
    ref = tcgen05.TMEMRef.from_alloc(addr_ref, shape, collective=c)
except ValueError as e:
    raise TypeError(f'bad tmem_addr_ref: {e}') from e

Prevention

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

When it happens

Trigger: 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.

Common situations: 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.

Understand the failure class

Background: Invalid argument type errors: "must be of type string", "expected X, got Y", and ERR_INVALID_ARG_TYPE explained — this error's family across 15 libraries.

Related errors

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/2864e625d16f3fa0. Report an issue: GitHub.