{"record":{"id":"821ea438e066bd12","repo":"jax-ml/jax","slug":"input-layouts-in-layouts-do-not-match-yield-layo","errorCode":null,"errorMessage":"Input layouts {in_layouts} do not match yield layouts {yield_layouts}","messagePattern":"Input layouts (.+?) do not match yield layouts (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2834,"sourceCode":"  after_block = while_op.after.blocks[0]\n  condition_op = before_block.operations[len(before_block.operations) - 1]\n  yield_op = after_block.operations[len(after_block.operations) - 1]\n\n  in_layouts = (\n      inference_utils.in_layouts(while_op)\n      if inference_utils.should_have_in_layout(while_op)\n      else []\n  )\n  out_layouts = (\n      inference_utils.out_layouts(while_op)\n      if inference_utils.should_have_out_layout(while_op)\n      else []\n  )\n\n  if in_layouts:\n    yield_layouts = inference_utils.in_layouts(yield_op)\n    if in_layouts != yield_layouts:\n      raise ValueError(\n          f\"Input layouts {in_layouts} do not match yield layouts\"\n          f\" {yield_layouts}\"\n      )\n\n  if out_layouts:\n    condition_layouts = inference_utils.in_layouts(condition_op)\n    if out_layouts != condition_layouts:\n      raise ValueError(\n          f\"Output layouts {out_layouts} do not match condition layouts\"\n          f\" {condition_layouts}\"\n      )\n\n  flat_inits, inits_template = _flatten_ir_values(while_op.inits, in_layouts)\n  result_types = _infer_flat_result_types(while_op, out_layouts)\n  new_while_op = scf.WhileOp(result_types, flat_inits)\n\n  # Before block\n  init_types = [v.type for v in flat_inits]","sourceCodeStart":2816,"sourceCodeEnd":2852,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2816-L2852","documentation":"For scf.while loops with layout-annotated inputs, the layouts of the loop's init/in values must equal the layouts of the values yielded by the body's yield op, for the same flattening reason as for-loops.","triggerScenarios":"scf.while loops where the body yields values whose inferred layouts differ from the incoming (init/arg) layouts of the carried values.","commonSituations":"While-loop-based iterative kernels that re-layout or broadcast carried accumulators inside the body.","solutions":["Ensure carried values keep the same layout across the while body; convert layouts before the loop","Compare inference_utils.in_layouts(yield_op) with in_layouts of the while op to locate the mismatched operand","Hoist layout-changing ops out of the loop body"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert inference_utils.in_layouts(while_op) == inference_utils.in_layouts(yield_op), 'while carry layouts must match yield'","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Treat while-loop carries as layout-invariant","Debug with inference_utils layout dumps when loops fail to lower"],"tags":["mosaic","gpu","scf","while-loop","layout-mismatch"],"backgroundTag":"layout-mismatch","analyzedSha":"1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb","analyzedAt":"2026-08-27T09:53:25.647Z","schemaVersion":2},"datasetVersion":"2026-08-27T13:17:12.746Z"}