jax-ml/jax · error · ValueError

Expected WGSplatFragLayout, but got {fa.layout}

Error message

Expected WGSplatFragLayout, but got {fa.layout}

What it means

Scalar outputs of inline_mgpu must use a WGSplatFragLayout fragment (a value splatted across the warp group). Any other fragment layout on a scalar-typed return raises ValueError.

Source

Thrown at jax/_src/pallas/mosaic_gpu/primitives.py:3850

        raise ValueError(f"Expected a FragmentedArray, but got: {fa}")
      if isinstance(result_ty, ir.VectorType):
        result_shape = ir.VectorType(result_ty).shape
        if fa.shape != tuple(result_shape):
          raise ValueError(f"Expected {result_shape} but got {fa.shape}")
        if out_layout != mgpu.layouts.to_layout_attr(fa.layout):
          raise ValueError(
              f"Output layout {out_layout} does not match the layout of the"
              f" returned fragmented array {fa.layout}."
          )
        ir_ret.append(
            mgpu.dialect_lowering.fragmented_array_to_ir(fa, result_ty)
        )
      else:  # scalar case.
        assert out_layout is None
        if fa.shape:
          raise ValueError(f"Expected 0D shape, but got {fa.shape}")
        if not isinstance(fa.layout, mgpu.WGSplatFragLayout):
          raise ValueError(f"Expected WGSplatFragLayout, but got {fa.layout}")
        value = fa.registers.item()
        ir_ret.append(value)

    mgpu.dialect.return_(ir_ret)


@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Warpgroup)
@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.WGxWARP_SEMANTICS)
def _inline_mgpu_lowering_rule_wg_semantics(
    ctx: lowering.LoweringRuleContext,
    *flat_args_and_transforms,
    mgpu_fn: Callable[..., Any],
    flat_arg_types,
    flat_ret_ty,
    pytree_args,
    pytree_ref_transforms,
    pytree_ret_ty,
):

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Broadcast or splat the scalar across the warp group before returning (use ops that produce WGSplatFragLayout)
2. Run under warp-group semantics where vector outputs are used instead of the scalar splat path

assert isinstance(fa.layout, mgpu.WGSplatFragLayout)

def is_wg_splat(fa): return isinstance(fa.layout, mgpu.WGSplatFragLayout)

When it happens

Trigger: Returning a scalar-typed FragmentedArray whose layout is not mgpu.WGSplatFragLayout in the scalar branch of the lowering.

Common situations: Producing the scalar via operations that yield a non-splat layout; mixing warp and warp-group produced values.

Related errors

• 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.
• Only TiledLayout supports reductions.

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