{"record":{"id":"0b09bd176381e1ca","repo":"jax-ml/jax","slug":"expected-a-fragmentedarray-but-got-fa","errorCode":null,"errorMessage":"Expected a FragmentedArray, but got: {fa}","messagePattern":"Expected a FragmentedArray, but got: (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3832,"sourceCode":"        )\n        fn_inputs.append(fa)\n      else:  # scalar case.\n        is_signed = mgpu_utils.is_signed(aval.dtype)\n        fa = mgpu.FragmentedArray.splat(arg, (), is_signed=is_signed)\n        fn_inputs.append(fa)\n\n    args = jax.tree.unflatten(pytree_args, fn_inputs)\n    inner_ret = mgpu_fn(ctx.launch_ctx, *args)\n    if inner_ret is None:\n      inner_ret = []\n    elif not isinstance(inner_ret, tuple) and not isinstance(inner_ret, list):\n      inner_ret = [inner_ret]\n    ir_ret = []\n    for fa, result_ty, out_layout in zip(\n        inner_ret, results_ty, out_layouts, strict=True\n    ):\n      if not isinstance(fa, mgpu.FragmentedArray):\n        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}\")","sourceCodeStart":3814,"sourceCodeEnd":3850,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3814-L3850","documentation":"Inside the inline_mgpu lowering, each returned value must be an mgpu.FragmentedArray when the corresponding result type is a vector type. Getting a plain array or other object raises this ValueError.","triggerScenarios":"An inline_mgpu callback under warp-group semantics returns a raw value (numpy array, scalar, etc.) where a FragmentedArray is expected by the lowering.","commonSituations":"Returning constants, host-side arrays, or scalars in vector-typed output slots; mismatch between declared vector return type and scalar-producing callback body.","solutions":["Ensure the callback produces FragmentedArray outputs (e.g. via mgpu ops that yield fragments)","If returning a scalar, declare the return type as scalar so the scalar path is taken"],"exampleFix":null,"handlingStrategy":"type-guard","validationCode":"from jax._src.pallas.mosaic_gpu import mgpu\nassert all(isinstance(v, mgpu.FragmentedArray) for v in rets)","typeGuard":"def is_fragmented(v): return isinstance(v, mgpu.FragmentedArray)","tryCatchPattern":null,"preventionTips":["Under warp-group semantics, produce outputs via mgpu ops yielding fragments"],"tags":["jax","pallas","fragmented-array","lowering"],"backgroundTag":"unsupported-return-type","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}