jax-ml/jax · error · ValueError

{ref} is not a memref.

Error message

{ref} is not a memref.

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2459

  output_shape = ir.MemRefType(op.tmem_ref.type).shape
  ncols = output_shape[1] // packing

  with utils.when(ctx.single_warp_per_block_predicate):
    tcgen05.tmem_dealloc(tmem_addr, ncols, collective, exact=False)

  return []


def _tmem_ref_from_ir(
    ref: ir.Value, expected_layout: ir.Attribute
) -> tcgen05.TMEMRef:
  """Returns a TMEMRef from an IR value.

  Throws an error if the annotated layout does not match the expected layout.
  """
  if not isinstance(ref.type, ir.MemRefType):
    raise ValueError(f"{ref} is not a memref.")
  mem_ref_ty = ir.MemRefType(ref.type)

  if mem_ref_ty.memory_space != utils.tmem():
    raise ValueError(
        f"{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM."
    )

  i32 = ir.IntegerType.get_signless(32)
  conversion_cast, [tmem_addr] = _undo_conversion_cast(ref, [i32])

  assert mem_ref_ty.rank == 2
  shape = cast(tuple[int, int], tuple(mem_ref_ty.shape))
  el_ty = mem_ref_ty.element_type
  layout_attr = conversion_cast.attributes["layout"]
  if layout_attr != expected_layout:
    raise ValueError(
        f"{ref} has a layout {layout_attr} that does not match the expected"
        f" layout {expected_layout}."

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Check the producer of the offending operand — it must be an mgpu.mem_ref / TMEM allocation producing a memref type
2. Ensure you wrapped the value with the proper TMEM allocation op (e.g. local_tensor with tmem memory space) before passing it
3. Verify no layout-inference step dropped the memref typing on the operand chain

Example fix

// before
acc = my_tensor  # tensor, not TMEM memref
mgpu.tcgen05_mma(a, b, acc)
// after
acc = mgpu.local_alloc(my_tensor, memory_space='tmem')
mgpu.tcgen05_mma(a, b, acc)

from jaxlib.mlir import ir
def is_memref(v):
    return isinstance(v.type, ir.MemRefType)

def is_tmem_operand(v) -> bool:
    from jaxlib.mlir import ir
    t = getattr(v, 'type', None)
    return isinstance(t, ir.MemRefType)

When it happens

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

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

Related errors

• {ref} has a memory space {mem_ref_ty.memory_space} that is n
• {ref} has a layout {layout_attr} that does not match the exp
• A custom return op must terminate the block.
• gpu.launch op not found.
• A scale layout {a_scale.layout} is not supported

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