{"record":{"id":"37eb5bd3bc00cb11","repo":"jax-ml/jax","slug":"inferred-layout-not-found-for-operand-operand","errorCode":null,"errorMessage":"Inferred layout not found for operand {operand}.","messagePattern":"Inferred layout not found for operand (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2695,"sourceCode":"  return layout1 != layout2 and not isinstance(layout1, fa.WGSplatFragLayout)\n\n\ndef check_for_expensive_relayout(module: ir.Module):\n  \"\"\"Returns whether the given module has an expensive relayout.\"\"\"\n  layout_for_variable: dict[ir.Value, fa.FragmentedLayout] = {}\n\n  def _check_for_expensive_relayout(op: ir.OpView):\n    if not inference_utils.should_have_layout(op):\n      return\n    if inference_utils.has_in_layouts_set(op):\n      in_layouts = iter(layouts_lib.from_layout_attr(l)\n                        for l in cast(ir.ArrayAttr, op.attributes[\"in_layouts\"]))\n      for operand in op.operands:\n        assert isinstance(operand, ir.Value)\n        if not isinstance(operand.type, ir.VectorType):\n          continue\n        if operand not in layout_for_variable:\n          raise ValueError(\n              f\"Inferred layout not found for operand {operand}.\"\n          )\n        operand_layout = next(in_layouts)\n        # TODO(bchetioui): refine to figure out whether it's a cheap relayout.\n        if _is_expensive_relayout(layout_for_variable[operand], operand_layout):\n          raise ValueError(\n              f\"Inferred layout {operand_layout} for operand {operand} does \"\n              f\"not match the layout in layout_for_variable \"\n              f\"{layout_for_variable[operand]}.\"\n          )\n    if inference_utils.has_out_layouts_set(op):\n      out_layouts = iter(layouts_lib.from_layout_attr(l)\n                         for l in cast(ir.ArrayAttr, op.attributes[\"out_layouts\"]))\n      for result in op.results:\n        assert isinstance(result, ir.Value)\n        if not isinstance(result.type, ir.VectorType):\n          continue\n        assert result not in layout_for_variable","sourceCodeStart":2677,"sourceCodeEnd":2713,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2677-L2713","documentation":"After solving the layout constraint system, every vector operand of an op with explicitly set in_layouts must have an inferred layout recorded in layout_for_variable. If a variable never got a layout assigned (e.g. it bypassed inference), this internal-consistency error fires.","triggerScenarios":"An op with in_layouts set consumes a vector value that was never registered in the constraint system (e.g. a constant or op excluded from layout inference), so lookup by ValueSite fails.","commonSituations":"Custom ops or manually inserted MLIR values feeding mgpu ops; edge cases with constants/splat materialization in Mosaic kernels.","solutions":["Ensure all values flowing into ops with explicit layouts are produced by ops that participate in layout inference","Rebuild the computation purely within the Mosaic DSL rather than mixing raw MLIR values","Update jax — missing-inferred-layout cases for constants are bug-fixed over time"],"exampleFix":null,"handlingStrategy":"try-catch","validationCode":null,"typeGuard":null,"tryCatchPattern":"try:\n    run_layout_inference(module)\nexcept ValueError as e:\n    if 'Inferred layout not found' in str(e):\n        # restructure kernel to avoid raw MLIR values feeding layout ops\n        raise","preventionTips":["Produce all operands via Mosaic DSL ops so they register in layout inference","Keep jax updated; these are frequently fixed bugs"],"tags":["jax","mosaic","layout-inference","internal"],"backgroundTag":"layout-inference-failure","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}