{"record":{"id":"58c85938bec3a668","repo":"jax-ml/jax","slug":"expected-the-same-number-of-out-attr-suffix-le","errorCode":null,"errorMessage":"Expected the same number of out_{attr_suffix} ({len(out_layouts)}) as {value_type} results ({num_matching_results}). op=\n  {op}","messagePattern":"Expected the same number of out_(.+?) \\((.+?)\\) as (.+?) results \\((.+?)\\)\\. op=\n  (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2319,"sourceCode":"    op: ir.OpView,\n) -> None:\n  \"\"\"Ensures that the right number of in/out layouts are provided for an op.\n\n  Layouts here are can be vector layouts, TMEM layouts, or SMEM transforms.\n  \"\"\"\n  layouts = lambda attr: op.attributes[attr] if attr in op.attributes else []\n  in_layouts = layouts(f\"in_{attr_suffix}\")\n  out_layouts = layouts(f\"out_{attr_suffix}\")\n\n  num_matching_operands = sum(map(filter_fn, op.operands))\n  if len(in_layouts) != num_matching_operands:\n    raise ValueError(\n        f\"Expected the same number of in_{attr_suffix} ({len(in_layouts)}) as \"\n        f\"{value_type} operands ({num_matching_operands}). op=\\n  {op}\"\n    )\n  num_matching_results = sum(map(filter_fn, op.results))\n  if len(out_layouts) != num_matching_results:\n    raise ValueError(\n        f\"Expected the same number of out_{attr_suffix} ({len(out_layouts)}) \"\n        f\"as {value_type} results ({num_matching_results}). op=\\n  {op}\"\n    )\n\n\n@dataclasses.dataclass(frozen=True)\nclass _TypeAndLayout:\n  type: ir.Type\n  layout: cs.Constant\n\n\ndef assign_layouts(solution: dict[ValueSite, cs.Constant]) -> None:\n  \"\"\"Assigns the layouts in `solution` to the MLIR ops they belong to.\n\n  This function requires that, for each MLIR op that appears in `solution`,\n  `solution` contains a layout assignment for all of its `vector`, TMEM, and\n  SMEM operands and results. Block arguments are ignored.\n  \"\"\"","sourceCodeStart":2301,"sourceCodeEnd":2337,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2301-L2337","documentation":"Same invariant check as for inputs, but for results: the number of out_layouts attributes must equal the number of vector-typed results of the op. A mismatch means the op's layout attributes were not built consistently with its results.","triggerScenarios":"Manually crafting an mgpu op whose out_layouts array length differs from its vector results count; ops produced by buggy lowering rules.","commonSituations":"Writing custom Mosaic lowering/layout rules, or hitting a regression after upgrading jax where result layout attr emission changed.","solutions":["Use the DSL wrappers (e.g. mgpu helpers or FraggedArray-returning APIs) that infer out layouts","Fix the custom op builder so one out layout is emitted per vector result","Report/check upstream if it reproduces with library-provided ops"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"n_res = sum(1 for r in op.results if isinstance(r.type, ir.VectorType))\nassert len(op.attributes.get('out_layouts', [])) == n_res","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Emit one out_layout per vector result in custom lowering rules","Run the module through Mosaic's verifier early in development"],"tags":["jax","mosaic","mlir","layout","invariant"],"backgroundTag":"invalid-operation-construction","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}