jax-ml/jax · error · TypeError

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

Error message

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

What it means

Warpgroup swap requires the destination to be an MLIR value of MemRefType (shared/global memory ref). A raw MemRefType or anything else is rejected, mirroring the lane-level type check for stores.

Source

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

    case _:
      raise NotImplementedError(f"Unsupported transforms: {transforms}")
  if ctx.module_ctx.auto_barriers:
    barrier()  # Make sure the writes have completed.
  return old_value


@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Warpgroup)
@register_lowering_rule(sp.swap_p, *gpu_core.WGxWARP_SEMANTICS)
def _swap_lowering_rule_wg(
    ctx: LoweringRuleContext, x_smem, value, *leaves, tree
):
  shape = ctx.avals_out[0].shape
  if shape and not isinstance(value.type, ir.VectorType):
    raise TypeError(f"Can only store scalars or vectors (got {value}).")
  if not (
      isinstance(x_smem, ir.Value) and isinstance(x_smem.type, ir.MemRefType)
  ):
    raise TypeError(f"Can only store to references (got {x_smem}).")
  if shape and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    raise NotImplementedError("Can only store scalars in warp-level lowering.")
  transforms = jax.tree.unflatten(tree, leaves)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[2:])
  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  x_smem, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_smem, transform_avals, transforms,
      allow_peer_refs=True
  )
  if transforms:
    raise NotImplementedError(
        "Transforms are not yet implemented for warpgroup semantics"
    )
  assert isinstance(x_smem, ir.Value)
  value = _ensure_ir_value(value, ctx.avals_in[1].dtype)
  if shape:
    old_value = mgpu.dialect.vector_load(x_smem)
    mgpu.dialect.vector_store(value, x_smem)

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Pass a proper materialized shared-memory reference
2. If writing custom lowering code, lower the ref through the standard path so it becomes ir.Value with MemRefType

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

When it happens

Trigger: Storing via swap_p under warpgroup semantics where x_smem is not a materialized memref value (internal misuse, peer refs, or custom lowering).

Common situations: Custom primitives or internal code paths passing unmaterialized refs; ref handling bugs after _handle_transforms.

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

• Can only store scalars or vectors (got {value}).
• ref must be a reference
• Unsupported type: {x}
• The base ref for aliases must come from a slice_smem op.
• Transforms are not yet implemented for warpgroup semantics

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