{"record":{"id":"b95cad646cbc251b","repo":"jax-ml/jax","slug":"output-layout-out-layout-does-not-match-the-layo","errorCode":null,"errorMessage":"Output layout {out_layout} does not match the layout of the returned fragmented array {fa.layout}.","messagePattern":"Output layout (.+?) does not match the layout of the returned fragmented array (.+?)\\.","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/_src/pallas/mosaic_gpu/primitives.py","lineNumber":3838,"sourceCode":"\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)\n\n","sourceCodeStart":3820,"sourceCodeEnd":3856,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/_src/pallas/mosaic_gpu/primitives.py#L3820-L3856","documentation":"For vector-typed outputs of inline_mgpu, the layout attribute computed from the declared output layout must equal the layout of the returned FragmentedArray. Divergent layouts (e.g. declared row-major but fragment is column-major) raise this ValueError.","triggerScenarios":"Returning a FragmentedArray whose fa.layout differs from the out_layout derived from the declared return type or layout hint.","commonSituations":"Operations inside the callback change the fragment layout (transpose, layout conversion) without a matching declared out_layout; stale layout hints.","solutions":["Set the declared output layout to match the fragment layout produced inside the callback","Insert an explicit layout conversion in the callback so the final layout matches the declaration"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert mgpu.layouts.to_layout_attr(fa.layout) == out_layout","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid layout-changing ops (transpose) at the tail of callbacks; convert layout explicitly"],"tags":["jax","pallas","layout-mismatch","mosaic-gpu"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}