jax-ml/jax · error · ValueError

Input layouts {in_layouts} do not match yield layouts {yield

Error message

Input layouts {in_layouts} do not match yield layouts {yield_layouts}

What it means

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.

Source

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

  after_block = while_op.after.blocks[0]
  condition_op = before_block.operations[len(before_block.operations) - 1]
  yield_op = after_block.operations[len(after_block.operations) - 1]

  in_layouts = (
      inference_utils.in_layouts(while_op)
      if inference_utils.should_have_in_layout(while_op)
      else []
  )
  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]

View on GitHub (pinned to 1e1c6a8fc0)

Solutions

1. Ensure carried values keep the same layout across the while body; convert layouts before the loop
2. Compare inference_utils.in_layouts(yield_op) with in_layouts of the while op to locate the mismatched operand
3. Hoist layout-changing ops out of the loop body

assert inference_utils.in_layouts(while_op) == inference_utils.in_layouts(yield_op), 'while carry layouts must match yield'

When it happens

Trigger: scf.while loops where the body yields values whose inferred layouts differ from the incoming (init/arg) layouts of the carried values.

Common situations: While-loop-based iterative kernels that re-layout or broadcast carried accumulators inside the body.

Related errors

• Output layouts {out_layouts} do not match condition layouts
• {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/821ea438e066bd12. Report an issue: GitHub.