{"record":{"id":"3198b6c4b843afd3","repo":"jax-ml/jax","slug":"inferred-layout-operand-layout-for-operand-oper","errorCode":null,"errorMessage":"Inferred layout {operand_layout} for operand {operand} does not match the layout in layout_for_variable {layout_for_variable[operand]}.","messagePattern":"Inferred layout (.+?) for operand (.+?) does not match the layout in layout_for_variable (.+?)\\.","errorType":"exception","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/layout_inference.py","lineNumber":2701,"sourceCode":"\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\n        layout_for_variable[result] = next(out_layouts)\n\n  for op in module.body:\n    traverse_op(op, _check_for_expensive_relayout)\n\n","sourceCodeStart":2683,"sourceCodeEnd":2719,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/layout_inference.py#L2683-L2719","documentation":"The layout inferred by the constraint system for an operand disagrees with the op's explicitly requested in_layout, and the relayout between them is expensive (not a cheap bitcast/reshape). Mosaic refuses to silently insert costly data movement.","triggerScenarios":"User sets a layout cast or in_layouts that conflicts with the layouts propagation naturally derives, e.g. feeding a value laid out for MMA into an op expecting a different non-compatible tiled layout.","commonSituations":"Explicit layout_cast / set layouts on ops in Mosaic kernels that fight the inferred layouts; mixing warps/registers layouts across op boundaries.","solutions":["Remove the explicit layout/cast and let inference pick layouts","Insert an explicit layout_cast where you accept the cost, instead of relying on implicit conversion","Restructure the kernel so producer and consumer agree on one layout (e.g. keep the MMA-compatible layout throughout)"],"exampleFix":"// before\nacc = mgpu.layout_cast(acc, some_other_layout)\nout = op_with_explicit_in_layout(acc)\n// after\nout = op_with_explicit_in_layout(acc)  # drop the cast; keep one layout","handlingStrategy":"validation","validationCode":"if _is_expensive_relayout(current_layout, requested_layout):\n    x = mgpu.layout_cast(x, requested_layout)  # make the move explicit","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Avoid pinning conflicting layouts on both producer and consumer","Chain incompatible layouts via explicit layout_cast calls"],"tags":["jax","mosaic","layout","conflict"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}