{"record":{"id":"d00c3ead6f0912c8","repo":"jax-ml/jax","slug":"expected-result-shape-but-got-fa-shape","errorCode":null,"errorMessage":"Expected {result_shape} but got {fa.shape}","messagePattern":"Expected (.+?) but got (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3836,"sourceCode":"        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}\")\n        value = fa.registers.item()\n        ir_ret.append(value)\n\n    mgpu.dialect.return_(ir_ret)","sourceCodeStart":3818,"sourceCodeEnd":3854,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3818-L3854","documentation":"The shape of the FragmentedArray returned from an inline_mgpu callback must equal the shape of the declared vector result type. A mismatch raises ValueError with both shapes.","triggerScenarios":"Callback returns a FragmentedArray whose .shape differs from ir.VectorType(result_ty).shape for that output.","commonSituations":"Declaring output avals with a different shape than what the callback computes; broadcasting applied inside the callback changing shapes.","solutions":["Make the declared return aval shape match exactly what the callback returns","Remove implicit reshapes or broadcasts in the callback, or declare the broadcast shape explicitly"],"exampleFix":"// before\nret_ty = jax.ShapeDtypeStruct((16, 16), jnp.float32)  # callback returns (8,8)\n// after\nret_ty = jax.ShapeDtypeStruct((8, 8), jnp.float32)","handlingStrategy":"validation","validationCode":"assert tuple(fa.shape) == tuple(ret_ty.shape), (fa.shape, ret_ty.shape)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Keep declared out shapes in sync with callback results; add asserts in tests"],"tags":["jax","pallas","shape-mismatch","lowering"],"backgroundTag":"shape-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}