jax-ml/jax · error · ValueError

tmem_addr_ref must contain a single element, got: {addr_ref_

Error message

tmem_addr_ref must contain a single element, got: {addr_ref_ty}

What it means

The TMEM base address occupies exactly one i32 word in shared memory; from_alloc requires math.prod(addr_ref_ty.shape) == 1 (a scalar/0-d memref, possibly with singleton dims). Multi-element memrefs are rejected because there is no defined way to pick the address among them.

Source

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

  @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??
    # warp_idx = utils.warp_idx(sync=False)
    # tmem_addr = arith.ori(tmem_addr, arith.shli(warp_idx, utils.c(21, i32)))
    return cls(tmem_addr, shape, dtype, layout)

  def slice(self, *idxs) -> TMEMRef:
    i32 = ir.IntegerType.get_signless(32)
    base_idx, slice_shape, is_squeezed = utils.parse_indices(idxs, self.shape)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Use a 0-d/scalar smem i32 memref for each allocation's address
2. Index the array memref first and pass a single-element subview (shape must reduce to product 1)
3. Allocate one address slot per tmem.alloc call

Example fix

# before
addr_ref = smem_alloca(i32, [4])  # memref<4xi32>
# after
addr_ref = smem_alloca(i32, [])  # single-element 0-d memref

import math
assert math.prod(ir.MemRefType(tmem_addr_ref.type).shape) == 1

def single_element_memref(ref) -> bool:
    t = getattr(ref, 'type', None)
    return isinstance(t, ir.MemRefType) and math.prod(t.shape) == 1

When it happens

Trigger: Passing a 1D memref like memref<4xi32> or memref<1x1xi32> with shape (1,1) is fine but memref<2xi32> is not; passing an array of alloc results.

Common situations: Allocating the address slot as a vector to 'be safe'; batching several allocations into one buffer and passing the whole buffer.

Related errors

• Minor dimension of shape must be divisible by packing, got:
• Unsupported shape: {shape}. TMEM references must have either
• tmem_addr_ref must be a memref or a pointer, got: {tmem_addr
• tmem_addr_ref must be an i32 memref, got: {addr_ref_ty}
• TMEM refs must have at least 32 rows, got: {shape[0]}

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