merge_decoders mixes up decoder_model and decoder_model_with_past

[xenova]

System Info

Optimum: main
Platform: Linux (Ubuntu)
Python version: 3.10.12

Who can help?

@JingyaHuang @michaelbenayoun @fxmarty

Information

• The official example scripts
• My own modified scripts

Tasks

• An officially supported task in the examples folder (such as GLUE/SQuAD, ...)
• My own task or dataset (give details below)

Reproduction (minimal, reproducible, runnable)

1. Following this tutorial, export a model with these two components

2. According to the docs and source code, the past model should be placed second:

optimum/optimum/onnx/graph_transformations.py

Lines 197 to 204 in 8b43dd2

	def merge_decoders(
	decoder: Union[ModelProto, Path, str],
	decoder_with_past: Union[ModelProto, Path, str],
	graph_name: str = "merged",
	producer_name: str = "optimum-onnx",
	save_path: Optional[Union[str, Path]] = None,
	strict: bool = True,
	) -> ModelProto:

So, run

from optimum.onnx.graph_transformations import merge_decoders
merge_decoders('./decoder_model.onnx', './decoder_with_past_model.onnx', strict=False)
# merge_decoders('./decoder_with_past_model.onnx', './decoder_model.onnx', strict=False) works

This results in this error:

ValueError: The second ModelProto should not have more outputs than the first.

3. Swap the arguments and watch it work.

This also affects the other tutorial here, which breaks if you set post-process to true.

Expected behavior

The docs/source code should be updated to be consistent, and/or the example code in the tutorial should work with post-process=True