jax-ml/jax · error · NotImplementedError

Only scalar memrefs are supported, got {ref_shape}

Error message

Only scalar memrefs are supported, got {ref_shape}

What it means

Mosaic GPU's layout inference treats memref.LoadOp and memref.StoreOp constraints as scalar accesses only: the referenced memref must have an empty shape or shape [1]. Loading from or storing to a memref with any larger shape raises NotImplementedError, since vector/tiled accesses must go through dedicated ops (async load/store, etc.).

Source

Thrown at jax/experimental/mosaic/gpu/layout_inference.py:2094

  return cs.ConstraintSystem(constraints=[cs.Equals(dest_var, collapse_expr)]), {
      source_var: [source],
      dest_var: [dest],
  }


# `memref.load` and `memref.store` are used to load barrier phases which are
# scalars---the rule needn't do anything interesting, but we need to have it.
@_add_constraint_system_derivation_rule(memref.LoadOp)
@_add_constraint_system_derivation_rule(memref.StoreOp)
def _memref_load_store_op_constraint_system(
    ctx: DerivationContext,
    op: memref.LoadOp | memref.StoreOp,
) -> ConstraintSystemDerivationRuleResult:
  del ctx

  ref_shape = ir.MemRefType(op.memref.type).shape
  if ref_shape and ref_shape != [1]:
    raise NotImplementedError(
        f"Only scalar memrefs are supported, got {ref_shape}"
    )

  ref_op_index = 0 if isinstance(op, memref.LoadOp) else 1
  ref = ValueSite(op, VariableType.OPERAND, ref_op_index)
  var = cs.Variable(ref)
  assignments: dict[cs.Variable, cs.Constant] = {var: cs.SMEMTransforms(None, None)}
  return cs.ConstraintSystem(assignments=assignments), {var: [ref]}


@_add_constraint_system_derivation_rule(mgpu.TryClusterCancelOp)
@_add_constraint_system_derivation_rule(mgpu.QueryClusterCancelOp)
def _cluster_launch_control_ops_constraint_system(
    ctx: DerivationContext,
    op: mgpu.TryClusterCancelOp | mgpu.QueryClusterCancelOp,
) -> ConstraintSystemDerivationRuleResult:
  ref = ValueSite(op, VariableType.OPERAND, 0)
  var = ctx.producer_ref(ref)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Index down to a scalar memref: ensure the memref type is scalar (shape [] or [1]) before memref.load/memref.store, e.g. by slicing/subviewing to single elements
2. Replace shaped load/store with Mosaic's supported tiled transfer ops (async_load/async_store, tma) which have their own constraint rules
3. If you meant element access, keep the memref scalar at allocation and compute addresses via offsets rather than a shaped buffer

Example fix

# before: elem = memref.load(buf, [i, j])  # buf: memref<128x64xf32>
elem = memref.load(buf[i, j], [])  # buf[i, j] is memref<f32> after subview

# or better: use tiled async copies
mgpu.async_load(source=gmem, destination=smem, ...)

shape = ir.MemRefType(op.memref.type).shape
if shape and shape != [1]:
    raise ValueError('use async_load/async_store for tiled access; load/store only on scalar memrefs')

def is_scalar_memref(v) -> bool:
    shape = ir.MemRefType(v.type).shape
    return not shape or shape == [1]

When it happens

Trigger: Emitting memref.load or memref.store whose memref operand has shape like [64], [128, 8], etc. — anything other than [] or [1] — in a Mosaic GPU kernel.

Common situations: Using element-wise load/store helpers on tiled shared-memory buffers instead of Mosaic's async copy primitives; code generated for older Mosaic versions where per-element loads on shaped memrefs with explicit indices were tolerated.

Related errors

• Only unit strides are supported but got {op.static_strides}.
• Transposed memrefs are not supported in ExpandShapeOp.
• CollapseShapeOp with empty reassociation is not supported.
• CollapseShapeOp with non-contiguous strides is not supported
• Scalars are not supported in async_store_smem

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