jax-ml/jax · error · ValueError

Expected a memref type but got {ref}

Error message

Expected a memref type but got {ref}

What it means

is_smem_ref (and its docstring) accept an ir.Value or ir.Type and check whether it is a MemRefType located in shared memory (memory_space == smem()). Anything that is not a memref — tensor, vector, function result types — raises this ValueError as a misuse of the predicate.

Source

Thrown at jax/experimental/mosaic/gpu/utils.py:2252

def tmem() -> ir.Attribute:
  """Returns the attribute for the TMEM memory space."""
  return ir.Attribute.parse("#mosaic_gpu.tmem")


def smem_cluster() -> ir.Attribute:
  """Returns the attribute for the cluster SMEM memory space."""
  return ir.Attribute.parse("#mosaic_gpu.smem_cluster")


def is_smem_ref(ref: ir.Value | ir.Type) -> bool:
  """Returns true if the input mem ref or memref type points to SMEM.

  If the input is not at all of a memref type, raises a ValueError.
  """
  if isinstance(ref, ir.Value):
    ref = ref.type
  if not isinstance(ref, ir.MemRefType):
    raise ValueError(f"Expected a memref type but got {ref}")
  ref = ir.MemRefType(ref)
  return ref.memory_space is not None and ref.memory_space == smem()


def is_tmem_ref(ref: ir.Value | ir.Type) -> bool:
  """Returns true if the input mem ref or memref type points to TMEM.

  If the input is not at all of a memref type, raises a ValueError.
  """
  if isinstance(ref, ir.Value):
    ref = ref.type
  if not isinstance(ref, ir.MemRefType):
    raise ValueError(f"Expected a memref type but got {ref}")
  ref = ir.MemRefType(ref)
  return ref.memory_space is not None and ref.memory_space == tmem()


def is_cluster_smem_ref(ref: ir.Value | ir.Type) -> bool:

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure the value has been lowered to buffers (memref) before calling is_smem_ref
2. Check isinstance(value.type, ir.MemRefType) first and handle the non-memref case explicitly
3. Fix the producer op to emit a memref result (e.g. use bufferization/to_memref)

Example fix

# before
in_smem = is_smem_ref(tensor_val)
# after
in_smem = isinstance(tensor_val.type, ir.MemRefType) and is_smem_ref(tensor_val)

if isinstance(ref, ir.Value):
    ref = ref.type
if isinstance(ref, ir.MemRefType):
    ok = is_smem_ref(ref)
else:
    ok = False  # or handle non-memref case

def is_memref(v_or_t) -> bool:
    t = v_or_t.type if isinstance(v_or_t, ir.Value) else v_or_t
    return isinstance(t, ir.MemRefType)

When it happens

Trigger: Calling is_smem_ref with a tensor-typed value, a vector value, or any non-memref ir.Type, e.g. checking the result of an op that produces a tensor.

Common situations: Feeding values from tensor-level (pre-lowering) IR into this memref-level helper; passing a module-level type or a function type accidentally; refactors where a value's type changed from memref to tensor.

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

• memref.cast transforms must have identical transforms for bo
• Unsupported dtype: {ref.dtype}
• Only byte-aligned shapes are supported. Got shape: {ref.dtyp
• No valid out swizzle{what}: minor dimension has {minor_dim_b
• Unsupported mesh type: {type(mesh)}

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