jax-ml/jax · error · NotImplementedError

Can only store scalars in warp-level lowering.

Error message

Can only store scalars in warp-level lowering.

What it means

In warp-level lowering of swap (store), only scalar values can be stored. Storing a shaped value under PrimitiveSemantics.Warp has no defined layout, so it raises NotImplementedError.

Source

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

    return mgpu.dialect.vector_load(x_ref, optimized=optimized)
  else:
    return memref_dialect.load(x_ref, [])


@register_lowering_rule(sp.swap_p, mgpu.LoweringSemantics.Lane)
@register_lowering_rule(sp.swap_p, *gpu_core.LANExWARP_SEMANTICS)
def _swap_lowering_rule(
    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:]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Store scalar elements (loop or broadcast per-lane scalars) inside warp semantics
2. Hoist the tensor store out to warpgroup/lane semantics
3. Verify the BlockSpec shape vs the semantics decorator so stores remain scalar

Example fix

# before
with plgpu.warp_semantics():
  x_ref[...] = vec  # vec has shape
# after
for i in range(vec.shape[0]):
  x_ref[i] = vec[i]  # scalar stores, or move store out of warp region

assert np.ndim(value_to_store) == 0 or semantics != 'warp'

When it happens

Trigger: Assigning a tensor (non-scalar) to a ref inside a plgpu.warp_semantics region: x_ref[...] = vec where avals_out[0].shape is non-empty under Warp semantics.

Common situations: Converting a lane-level store loop into warp-level code without making stores per-scalar; using vector ops inside warp semantics.

Related errors

• packed cannot be specified if layout is specified.
• packed, collective and layout arguments are only supported f
• Cannot specify both out_shape and out_type. Use out_type.
• Cannot specify both scratch_shapes and scratch_types. Use sc
• Unsupported dtype: {ref.dtype}

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