in optimum/exporters/onnx/base.py [0:0]
def outputs(self) -> Dict[str, Dict[int, str]]:
common_outputs = super(OnnxConfigWithPast, self).outputs
# Renaming the outputs axes properly.
for name, axes_names in common_outputs.items():
if self._behavior is ConfigBehavior.ENCODER or "encoder" in name:
sequence_name = "encoder_sequence_length"
else:
sequence_name = "decoder_sequence_length"
new_axes_names = {}
for axis_idx, axis_name in axes_names.items():
if "sequence" in axis_name:
if self.use_past_in_inputs is False or self.is_merged is True:
new_axes_names[axis_idx] = sequence_name
else:
# Trick to force it since ONNX sometimes infer a dynamic axis where it's not.
new_axes_names[axis_idx] = "1"
else:
new_axes_names[axis_idx] = axis_name
common_outputs[name] = new_axes_names
if self.use_past:
# When exporting decoder models with use_cache=True, both the decoder without past and with past have the KV cache as an output.
self.add_past_key_values(common_outputs, direction="outputs")
return common_outputs