jax-ml/jax · error · ValueError

Only plgpu.Layout.WG_SPLAT is supported for scalar values.

Error message

Only plgpu.Layout.WG_SPLAT is supported for scalar values.

What it means

When relayout_p is lowered at Warpgroup semantics on a 0-d (scalar) value, the only supported target layout is WGSplatFragLayout. Requesting any other layout for a scalar raises this ValueError.

Source

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

      "relayout_p is not supported with Lane semantics."
  )


@register_lowering_rule(pjit.relayout_p, mgpu.LoweringSemantics.Warpgroup)
def _relayout_lowering_wg(
    ctx: LoweringRuleContext, x, *, dst_layout
):
  if dst_layout is jax_layout.AutoLayout:
    return x
  layout = fa.TiledLayout(
      dst_layout.tiling,
      dst_layout.warp_dims,
      dst_layout.lane_dims,
      dst_layout.vector_dim,
  )
  if ctx.avals_in[0].ndim == 0:  # scalar case
    if layout != mgpu.WGSplatFragLayout():
      raise ValueError(
          "Only plgpu.Layout.WG_SPLAT is supported for scalar values."
      )
    return x
  return mgpu.dialect.layout_cast(x, mgpu.to_layout_attr(layout))


@register_lowering_rule(gpu_core.layout_cast_p, mgpu.LoweringSemantics.Lane)
def _layout_cast_lowering(ctx: LoweringRuleContext, x, *, new_layout):
  del ctx  # Unused.
  return x.to_layout(new_layout.to_mgpu())


@register_lowering_rule(gpu_core.layout_cast_p, mgpu.LoweringSemantics.Warpgroup)
def _layout_cast_lowering_wg(
    ctx: LoweringRuleContext, x, *, new_layout
):
  layout = new_layout.to_mgpu()
  if ctx.avals_in[0].ndim == 0:  # scalar case

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Skip layout_cast for scalars (only cast arrays with ndim>0)
2. Use plgpu.Layout.WG_SPLAT for scalar values

Example fix

# before
x = plgpu.layout_cast(scalar_val, tiled_layout)
# after
x = scalar_val  # scalars need no cast; or use WG_SPLAT

if x.ndim == 0: assert layout == plgpu.Layout.WG_SPLAT

def needs_layout_cast(x) -> bool:
    return getattr(x, 'ndim', 0) > 0

When it happens

Trigger: Calling lax.layout_cast / relayout with a non-WG_SPLAT layout on a scalar inside a Warpgroup-semantic Mosaic GPU kernel.

Common situations: Generic code that applies a fixed tiling layout to all values, including scalars like loop counters or accumulated scalars.

Related errors

• Expected WGSplatFragLayout, but got {fa.layout}
• packed cannot be specified if layout is specified.
• Cannot commute `UntilingTransform` with `ReshapeTransform` w
• Swizzle {self.swizzle} requires the trailing dimension to be
• Can't instantiate {self} with arguments.

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