jax-ml/jax · error · TypeError

Can only store scalars or vectors

Error message

Can only store scalars or vectors (got {value}).

What it means

Warpgroup swap lowering requires the stored value to be a scalar (no shape) or an MLIR vector. A shaped value that is not an ir.VectorType (e.g. a FragmentedArray or a non-vector lowered value) is a TypeError.

Solutions

1. Ensure values stored under warpgroup semantics are lowered as vectors or scalars
2. Use the standard library ops (plgpu ops) rather than hand-building values for the store
3. File/check against recent JAX if using only public Pallas APIs

import jax._src.interpreters.mlir as mlir
def is_scalar_or_vector(v) -> bool:
    return not getattr(v, 'type', None) and True or isinstance(getattr(v, 'type', None), mlir.ir.VectorType)

When it happens

Trigger: Calling x_smem[...] = value under warpgroup semantics where value has shape but its lowered type is not ir.VectorType — usually an internal lowering inconsistency or a custom primitive feeding the wrong IR type.

Common situations: Writing custom Pallas primitives/lowering rules; mixing lane-semantics FragmentedArray values into warpgroup stores.

Understand the failure class

Background: "Must be a positive integer", "Invalid value", "Unsupported": the invalid-argument-value error family, when a library rejects the value you pass — this error's family across 35 libraries.

Related errors

• Can only store to references
• Can only store to references
• ref must be a reference
• The base ref for aliases must come from a slice_smem op.
• Transfer is not a multiple of

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

Appendix: source

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

          old_value = mgpu.FragmentedArray.load_strided(
              x_smem, is_signed=mgpu_utils.is_signed(v_aval.dtype)
          )
          value.store_untiled(x_smem)
    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)

View on GitHub (pinned to 1e1c6a8fc0)