{"record":{"id":"f84e0d678f7bae6c","repo":"jax-ml/jax","slug":"output-layouts-out-layouts-do-not-match-conditio","errorCode":null,"errorMessage":"Output layouts {out_layouts} do not match condition layouts {condition_layouts}","messagePattern":"Output layouts (.+?) do not match condition layouts (.+?)","errorType":"validation","errorClass":"ValueError","httpStatus":null,"severity":"error","filePath":"jax/experimental/mosaic/gpu/dialect_lowering.py","lineNumber":2842,"sourceCode":"  )\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]\n  new_before_block = new_while_op.before.blocks.append(*init_types)\n  results_template = _move_scf_block_to_block_with_flattened_arguments(\n      ctx,\n      before_block,\n      new_before_block,\n      scf.ConditionOp,\n      inits_template,\n  )","sourceCodeStart":2824,"sourceCodeEnd":2860,"githubUrl":"https://github.com/jax-ml/jax/blob/1e1c6a8fc06dfcd1247076ec5cae4640cea5d7bb/jax/experimental/mosaic/gpu/dialect_lowering.py#L2824-L2860","documentation":"For scf.while loops with layout-annotated outputs, the after-block's condition op must pass carry values whose layouts match the while op's result layouts; otherwise flattening cannot proceed.","triggerScenarios":"A while loop whose condition op forwards re-laid-out or differently-tiled values compared to the loop's declared out_layouts.","commonSituations":"While-loop kernels where the condition block yields a value produced by a layout-changing op (transpose_layout, slice with new layout).","solutions":["Keep condition-op args in the same layout as the while results","Move layout conversions before the while loop or restructure so carried layout is invariant","Inspect out_layouts vs condition_layouts operand-by-operand to find the offender"],"exampleFix":null,"handlingStrategy":"validation","validationCode":"assert inference_utils.out_layouts(while_op) == inference_utils.in_layouts(condition_op)","typeGuard":null,"tryCatchPattern":null,"preventionTips":["Condition op must forward values with the same layouts as while results","Avoid layout transforms inside while bodies"],"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"}