jax-ml/jax · error · TypeError

Can only store to references (got {x_ref}).

Error message

Can only store to references (got {x_ref}).

What it means

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.

Source

Thrown at jax/_src/pallas/mosaic_gpu/lowering.py:2272

    ctx: LoweringRuleContext, x_ref, value, *leaves, tree
):
  if isinstance(x_ref, tcgen05.TMEMRef):
    raise RuntimeError(
        "Stores to TMEM are asynchronous operations and cannot be performed"
        " using the usual syntax. Please use plgpu.async_store_tmem instead."
    )
  barrier = mgpu.warpgroup_barrier
  if ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    if ctx.avals_out[0].shape:
      raise NotImplementedError("Can only store scalars in warp-level lowering.")
    i32 = ir.IntegerType.get_signless(32)
    barrier = functools.partial(
        nvvm_dialect.bar_warp_sync, arith_dialect.constant(i32, -1)
    )
  value = _ensure_fa(value, ctx.avals_in[1].dtype)

  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only store to references (got {x_ref}).")
  v_aval = ctx.avals_in[1]
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])

  if ctx.module_ctx.auto_barriers:
    barrier()  # Make sure reads have completed before we write.

  if transforms and isinstance(transforms[0], gpu_core.UnswizzleRef):
    swizzle = transforms[0].swizzle
    transforms = transforms[1:]
    transform_avals = transform_avals[1:]
  elif (
      len(transforms) > 1
      and isinstance(transforms[0], gpu_core.ExtractAliasedRef)
      and isinstance(transforms[1], gpu_core.UnswizzleRef)
  ):
    swizzle = transforms[1].swizzle
    transforms = [transforms[0], *transforms[2:]]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass the Pallas reference object (Ref) into the store, not a lowered memref type
2. If writing custom lowering code, ensure refs are materialized to ir.Value via the standard ref-lowering path
3. Check for accidental unpacking of the ref into its type before the store

def is_valid_store_target(x) -> bool:
    import jax._src.interpreters.mlir as mlir
    return isinstance(x, mlir.ir.Value) and not isinstance(x, mlir.ir.MemRefType)

When it happens

Trigger: 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.

Common situations: Internal library misuse, custom primitives calling the lowering API directly, or passing an already-lowered memref type where a reference is expected.

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

• ref must be a reference
• Unsupported type: {x}
• Can only store scalars or vectors (got {value}).
• Can only store to references (got {x_smem}).
• Unknown action: {action}

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