jax-ml/jax · error · TypeError

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

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.

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)

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: 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 to references (got {x_smem}).
• 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/fb4012dd0a2fffa4. Report an issue: GitHub.