jax-ml/jax · error · ValueError

Output layouts {out_layouts} do not match condition layouts

Error message

Output layouts {out_layouts} do not match condition layouts {condition_layouts}

What it means

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.

Source

Thrown at jax/experimental/mosaic/gpu/dialect_lowering.py:2842

  )
  out_layouts = (
      inference_utils.out_layouts(while_op)
      if inference_utils.should_have_out_layout(while_op)
      else []
  )

  if in_layouts:
    yield_layouts = inference_utils.in_layouts(yield_op)
    if in_layouts != yield_layouts:
      raise ValueError(
          f"Input layouts {in_layouts} do not match yield layouts"
          f" {yield_layouts}"
      )

  if out_layouts:
    condition_layouts = inference_utils.in_layouts(condition_op)
    if out_layouts != condition_layouts:
      raise ValueError(
          f"Output layouts {out_layouts} do not match condition layouts"
          f" {condition_layouts}"
      )

  flat_inits, inits_template = _flatten_ir_values(while_op.inits, in_layouts)
  result_types = _infer_flat_result_types(while_op, out_layouts)
  new_while_op = scf.WhileOp(result_types, flat_inits)

  # Before block
  init_types = [v.type for v in flat_inits]
  new_before_block = new_while_op.before.blocks.append(*init_types)
  results_template = _move_scf_block_to_block_with_flattened_arguments(
      ctx,
      before_block,
      new_before_block,
      scf.ConditionOp,
      inits_template,
  )

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Keep condition-op args in the same layout as the while results
2. Move layout conversions before the while loop or restructure so carried layout is invariant
3. Inspect out_layouts vs condition_layouts operand-by-operand to find the offender

assert inference_utils.out_layouts(while_op) == inference_utils.in_layouts(condition_op)

When it happens

Trigger: A while loop whose condition op forwards re-laid-out or differently-tiled values compared to the loop's declared out_layouts.

Common situations: While-loop kernels where the condition block yields a value produced by a layout-changing op (transpose_layout, slice with new layout).

Related errors

• Input layouts {in_layouts} do not match yield layouts {yield
• {ref} has a layout {layout_attr} that does not match the exp
• Unsupported op type: {op}
• Rule cannot handle an op with vector operands or results: {o
• Mosaic GPU does not yet support AMD ROCm devices. Use ``comp

AI-assisted analysis of jax-ml/jax@1e1c6a8fc0 (2026-08-27). Data as JSON: /api/errors/f84e0d678f7bae6c. Report an issue: GitHub.