jax-ml/jax · error · ValueError

{ref} has a memory space {mem_ref_ty.memory_space} that is n

Error message

{ref} has a memory space {mem_ref_ty.memory_space} that is not TMEM.

What it means

Mosaic GPU's TMEM lowering requires the operand memref's memory_space to equal utils.tmem(). This error means the value is a valid memref but lives in a different memory space (e.g. shared or global), so it cannot be treated as a tcgen05 TMEM reference.

Source

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

  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}."
    )
  layout = layouts_lib.from_layout_attr(layout_attr)
  assert isinstance(layout, fa.TiledLayout)
  tmem_layout = tcgen05.TMEMLayout(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Allocate the buffer with the TMEM memory space (e.g. mgpu.local_tensor(..., memory_space=utils.tmem()))
2. For data in smem use the async_copy_smem_to_tmem lowering path instead of feeding smem memrefs directly
3. Check any custom layout-cast / slice producers preserve the tmem memory space attribute

Example fix

// before
acc = mgpu.local_tensor(shape, dtype, memory_space='shared')
// after
from jax.experimental.mosaic.gpu import utils
acc = mgpu.local_tensor(shape, dtype, memory_space=utils.tmem())

from jaxlib.mlir import ir
from jax.experimental.mosaic.gpu import utils
def in_tmem(v) -> bool:
    t = v.type
    return isinstance(t, ir.MemRefType) and t.memory_space == utils.tmem()

def is_tmem_ref(v) -> bool:
    from jaxlib.mlir import ir
    return isinstance(v.type, ir.MemRefType) and str(v.type.memory_space) == '#gpu.memory_space<tmem>'

When it happens

Trigger: Passing an smem local_tensor or a global-tensor-derived memref into TMEM operands of tcgen05_mma / async_copy_smem_to_tmem / slice_tmem etc., where memory_space != tmem.

Common situations: Forgetting to set memory_space='tmem' on the allocation; using a shared-memory scratch buffer directly as an accumulator in a Blackwell MMA pipeline instead of copying it to TMEM first.

Related errors

• {ref} is not a memref.
• Unsupported memory space: {orig_ref_ty.memory_space}
• Unsupported memory space: {ref_type.memory_space}
• {ref} has a layout {layout_attr} that does not match the exp
• A custom return op must terminate the block.

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