{"record":{"id":"cccdadf1c20c923e","repo":"jax-ml/jax","slug":"only-scalar-memrefs-are-supported-got-ref-shape","errorCode":null,"errorMessage":"Only scalar memrefs are supported, got {ref_shape}","messagePattern":"Only scalar memrefs are supported, got (.+?)","errorType":"validation","errorClass":"NotImplementedError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2094,"sourceCode":"  return cs.ConstraintSystem(constraints=[cs.Equals(dest_var, collapse_expr)]), {\n      source_var: [source],\n      dest_var: [dest],\n  }\n\n\n# `memref.load` and `memref.store` are used to load barrier phases which are\n# scalars---the rule needn't do anything interesting, but we need to have it.\n@_add_constraint_system_derivation_rule(memref.LoadOp)\n@_add_constraint_system_derivation_rule(memref.StoreOp)\ndef _memref_load_store_op_constraint_system(\n    ctx: DerivationContext,\n    op: memref.LoadOp | memref.StoreOp,\n) -> ConstraintSystemDerivationRuleResult:\n  del ctx\n\n  ref_shape = ir.MemRefType(op.memref.type).shape\n  if ref_shape and ref_shape != [1]:\n    raise NotImplementedError(\n        f\"Only scalar memrefs are supported, got {ref_shape}\"\n    )\n\n  ref_op_index = 0 if isinstance(op, memref.LoadOp) else 1\n  ref = ValueSite(op, VariableType.OPERAND, ref_op_index)\n  var = cs.Variable(ref)\n  assignments: dict[cs.Variable, cs.Constant] = {var: cs.SMEMTransforms(None, None)}\n  return cs.ConstraintSystem(assignments=assignments), {var: [ref]}\n\n\n@_add_constraint_system_derivation_rule(mgpu.TryClusterCancelOp)\n@_add_constraint_system_derivation_rule(mgpu.QueryClusterCancelOp)\ndef _cluster_launch_control_ops_constraint_system(\n    ctx: DerivationContext,\n    op: mgpu.TryClusterCancelOp | mgpu.QueryClusterCancelOp,\n) -> ConstraintSystemDerivationRuleResult:\n  ref = ValueSite(op, VariableType.OPERAND, 0)\n  var = ctx.producer_ref(ref)","sourceCodeStart":2076,"sourceCodeEnd":2112,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2076-L2112","documentation":"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.).","triggerScenarios":"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.","commonSituations":"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.","solutions":["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","Replace shaped load/store with Mosaic's supported tiled transfer ops (async_load/async_store, tma) which have their own constraint rules","If you meant element access, keep the memref scalar at allocation and compute addresses via offsets rather than a shaped buffer"],"exampleFix":"# before: elem = memref.load(buf, [i, j])  # buf: memref<128x64xf32>\nelem = memref.load(buf[i, j], [])  # buf[i, j] is memref<f32> after subview\n\n# or better: use tiled async copies\nmgpu.async_load(source=gmem, destination=smem, ...)","handlingStrategy":"validation","validationCode":"shape = ir.MemRefType(op.memref.type).shape\nif shape and shape != [1]:\n    raise ValueError('use async_load/async_store for tiled access; load/store only on scalar memrefs')","typeGuard":"def is_scalar_memref(v) -> bool:\n    shape = ir.MemRefType(v.type).shape\n    return not shape or shape == [1]","tryCatchPattern":null,"preventionTips":["Use Mosaic async copy primitives for any shaped buffer transfer","Reserve memref.load/store for scalar scratch slots (accumulators, flags)"],"tags":["jax","mosaic-gpu","memref","load-store","layout-inference","scalar"],"backgroundTag":"unsupported-operation-parameter","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}