{"record":{"id":"812071bdce7de719","repo":"jax-ml/jax","slug":"expected-the-same-number-of-in-attr-suffix-len","errorCode":null,"errorMessage":"Expected the same number of in_{attr_suffix} ({len(in_layouts)}) as {value_type} operands ({num_matching_operands}). op=\n  {op}","messagePattern":"Expected the same number of in_(.+?) \\((.+?)\\) as (.+?) operands \\((.+?)\\)\\. op=\n  (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2313,"sourceCode":"\n\ndef _ensure_right_number_of_layouts(\n    filter_fn: Callable[[ir.Value], bool],\n    attr_suffix: str,\n    value_type: str,\n    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:","sourceCodeStart":2295,"sourceCodeEnd":2331,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2295-L2331","documentation":"During layout inference finalization, Mosaic checks that the number of in_layouts/in_vectors/in_indices attributes on an op equals the number of vector-typed operands. A mismatch means the op was constructed with missing or extra layout attributes, which is an invariant violation of the DSL.","triggerScenarios":"Building an mgpu dialect op manually (via ir insertion or by bypassing the Python DSL wrappers) with fewer/more in_layouts entries than vector operands; or a bug in Mosaic's own op construction for a new op type.","commonSituations":"Extending Mosaic with new custom ops, mixing MLIR-level op building with the Python DSL, or version skew after layout attr naming changes (in_layouts vs in_vectors).","solutions":["Construct the op through the public mgpu Python helpers instead of raw MLIR op creation so layouts are attached automatically","Check for duplicate/missing entries in the in_* attribute array and make it match the count of vector operands","Update JAX/Mosaic to the latest version in case the op signature changed"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"n_vec = sum(1 for o in op.operands if isinstance(o.type, ir.VectorType))\nassert len(op.attributes.get('in_layouts', [])) == n_vec","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Prefer DSL helpers over raw MLIR op construction","Add assertions on attr-count vs operand-count in custom op builders"],"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"}