{"record":{"id":"57c6b7b01d7a5951","repo":"jax-ml/jax","slug":"expected-wgsplatfraglayout-but-got-fa-layout","errorCode":null,"errorMessage":"Expected WGSplatFragLayout, but got {fa.layout}","messagePattern":"Expected WGSplatFragLayout, but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3850,"sourceCode":"        raise ValueError(f\"Expected a FragmentedArray, but got: {fa}\")\n      if isinstance(result_ty, ir.VectorType):\n        result_shape = ir.VectorType(result_ty).shape\n        if fa.shape != tuple(result_shape):\n          raise ValueError(f\"Expected {result_shape} but got {fa.shape}\")\n        if out_layout != mgpu.layouts.to_layout_attr(fa.layout):\n          raise ValueError(\n              f\"Output layout {out_layout} does not match the layout of the\"\n              f\" returned fragmented array {fa.layout}.\"\n          )\n        ir_ret.append(\n            mgpu.dialect_lowering.fragmented_array_to_ir(fa, result_ty)\n        )\n      else:  # scalar case.\n        assert out_layout is None\n        if fa.shape:\n          raise ValueError(f\"Expected 0D shape, but got {fa.shape}\")\n        if not isinstance(fa.layout, mgpu.WGSplatFragLayout):\n          raise ValueError(f\"Expected WGSplatFragLayout, but got {fa.layout}\")\n        value = fa.registers.item()\n        ir_ret.append(value)\n\n    mgpu.dialect.return_(ir_ret)\n\n\n@lowering.register_lowering_rule(inline_mgpu_p, mgpu.LoweringSemantics.Warpgroup)\n@lowering.register_lowering_rule(inline_mgpu_p, *gpu_core.WGxWARP_SEMANTICS)\ndef _inline_mgpu_lowering_rule_wg_semantics(\n    ctx: lowering.LoweringRuleContext,\n    *flat_args_and_transforms,\n    mgpu_fn: Callable[..., Any],\n    flat_arg_types,\n    flat_ret_ty,\n    pytree_args,\n    pytree_ref_transforms,\n    pytree_ret_ty,\n):","sourceCodeStart":3832,"sourceCodeEnd":3868,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3832-L3868","documentation":"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.","triggerScenarios":"Returning a scalar-typed FragmentedArray whose layout is not mgpu.WGSplatFragLayout in the scalar branch of the lowering.","commonSituations":"Producing the scalar via operations that yield a non-splat layout; mixing warp and warp-group produced values.","solutions":["Broadcast or splat the scalar across the warp group before returning (use ops that produce WGSplatFragLayout)","Run under warp-group semantics where vector outputs are used instead of the scalar splat path"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert isinstance(fa.layout, mgpu.WGSplatFragLayout)","typeGuard":"def is_wg_splat(fa): return isinstance(fa.layout, mgpu.WGSplatFragLayout)","tryCatchPattern":null,"preventionTips":["Splat scalars across the warp group before returning them"],"tags":["jax","pallas","layout","scalar","mosaic-gpu"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}