{"record":{"id":"da4a0cc248d80ff7","repo":"jax-ml/jax","slug":"can-only-store-to-references-got-x-ref","errorCode":null,"errorMessage":"Can only store to references (got {x_ref}).","messagePattern":"Can only store to references \\(got (.+?)\\)\\.","errorType":"exception","errorClass":"TypeError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/lowering.py","lineNumber":2272,"sourceCode":"    ctx: LoweringRuleContext, x_ref, value, *leaves, tree\n):\n  if isinstance(x_ref, tcgen05.TMEMRef):\n    raise RuntimeError(\n        \"Stores to TMEM are asynchronous operations and cannot be performed\"\n        \" using the usual syntax. Please use plgpu.async_store_tmem instead.\"\n    )\n  barrier = mgpu.warpgroup_barrier\n  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:\n    if ctx.avals_out[0].shape:\n      raise NotImplementedError(\"Can only store scalars in warp-level lowering.\")\n    i32 = ir.IntegerType.get_signless(32)\n    barrier = functools.partial(\n        nvvm_dialect.bar_warp_sync, arith_dialect.constant(i32, -1)\n    )\n  value = _ensure_fa(value, ctx.avals_in[1].dtype)\n\n  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):\n    raise TypeError(f\"Can only store to references (got {x_ref}).\")\n  v_aval = ctx.avals_in[1]\n  transforms = jax.tree.unflatten(tree, leaves)\n  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])\n\n  if ctx.module_ctx.auto_barriers:\n    barrier()  # Make sure reads have completed before we write.\n\n  if transforms and isinstance(transforms[0], gpu_core.UnswizzleRef):\n    swizzle = transforms[0].swizzle\n    transforms = transforms[1:]\n    transform_avals = transform_avals[1:]\n  elif (\n      len(transforms) > 1\n      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)\n      and isinstance(transforms[1], gpu_core.UnswizzleRef)\n  ):\n    swizzle = transforms[1].swizzle\n    transforms = [transforms[0], *transforms[2:]]","sourceCodeStart":2254,"sourceCodeEnd":2290,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/lowering.py#L2254-L2290","documentation":"Type-check in the lane-level swap lowering rule: the destination must be a lowered reference (an MLIR value of MemRefType). Getting a raw MLIR MemRefType instead of a value means a non-reference (e.g. an unmaterialized memref type) was passed as the store target.","triggerScenarios":"Calling swap_p (x_ref[...] = v) where x_ref lowered to an ir.MemRefType object rather than an ir.Value with memref type — typically an internal/lower-level misuse or a peer ref not materialized.","commonSituations":"Internal library misuse, custom primitives calling the lowering API directly, or passing an already-lowered memref type where a reference is expected.","solutions":["Pass the Pallas reference object (Ref) into the store, not a lowered memref type","If writing custom lowering code, ensure refs are materialized to ir.Value via the standard ref-lowering path","Check for accidental unpacking of the ref into its type before the store"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":null,"typeGuard":"def is_valid_store_target(x) -> bool:\n    import jax._src.interpreters.mlir as mlir\n    return isinstance(x, mlir.ir.Value) and not isinstance(x, mlir.ir.MemRefType)","tryCatchPattern":null,"preventionTips":["Only public Pallas ref objects should appear as store targets","Avoid custom lowering code passing raw MemRefType"],"tags":["jax","pallas","mosaic-gpu","type-error","internal-api"],"backgroundTag":"invalid-argument-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}