jax-ml/jax · error · NotImplementedError

Transforms are not yet implemented for warpgroup semantics

Error message

Transforms are not yet implemented for warpgroup semantics

What it means

The warpgroup-semantics load lowering rule rejects any reference transforms (slicing/tiling/swizzle descriptors) on the loaded reference. Warpgroup lowering only supports plain shared/global memory loads without transform chains.

Source

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

def _get_lowering_rule_wg(
    ctx: LoweringRuleContext, x_ref, *leaves, tree, optimized=True
):
  if not isinstance(x_ref, ir.Value) and isinstance(x_ref, ir.MemRefType):
    raise TypeError(f"Can only load from references (got {x_ref}).")
  shape = ctx.avals_out[0].shape
  if shape and ctx.module_ctx.primitive_semantics == gpu_core.PrimitiveSemantics.Warp:
    raise ValueError("Can only load scalars in warp-level code.")

  transforms = jax.tree.unflatten(tree, leaves)
  assert isinstance(ctx.avals_in[0], state_types.AbstractRef)
  transform_avals = jax.tree.unflatten(tree, ctx.avals_in[1:])
  x_ref, _, transforms = _handle_transforms(
      ctx, ctx.avals_in[0], x_ref, transform_avals, transforms,
      allow_peer_refs=True
  )

  if transforms:
    raise NotImplementedError(
        "Transforms are not yet implemented for warpgroup semantics"
    )

  assert isinstance(x_ref, ir.Value)
  shape = ctx.avals_out[0].shape
  if shape:
    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(

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Apply the slicing to the underlying memory before entering warpgroup semantics, or use a separately allocated ref
2. Express the operation with explicit TMEM/shared-memory loads (e.g. mgpu/load from tmem) that don't require transforms
3. Restructure so the value is loaded without any transform chain

Example fix

# before
with plgpu.warp_group_semantics():
  v = x_ref[0:128, :]  # transform under warpgroup
# after
v = x_ref[0:128, :]  # slice at lane semantics / before entering wg region

When it happens

Trigger: Passing a transformed ref (e.g. x_ref[2:5, :] or a swizzled/tilted ref via plgpu transform descriptors) into a load lowered with LoweringSemantics.Warpgroup.

Common situations: Reusing lane-level kernel code that indexes into refs inside a warpgroup-scoped region (e.g. around tcgen05 MMA blocks); dynamic_slice or slicing on a ref under warpgroup semantics.

Related errors

• Transpose cannot be moved before a tiling transform when it
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Commuting a `UntilingTransform` with a `ReshapeTransform` is
• Unsupported transform: {type(transform)}
• The base ref for aliases must come from a slice_smem op.

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