AndroidRemoteController/python/py/Lib/site-packages/onnxruntime/transformers/convert_generation.py

# -------------------------------------------------------------------------
# Copyright (c) Microsoft Corporation. All rights reserved.
# Licensed under the MIT License.
# -------------------------------------------------------------------------
"""
This converts GPT2 or T5 model to onnx with beam search operator.

Example 1: convert gpt2 model with beam search:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx

Example 2: convert gpt2 model with beam search containing specific cuda optimizations:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx --use_gpu               \
        --past_present_share_buffer --use_decoder_masked_attention

Example 3: convert gpt2 model with beam search with mixed precision and enable SkipLayerNorm strict mode:
    python convert_generation.py -m gpt2 --output gpt2_beam_search.onnx --use_gpu -p fp16 --use_sln_strict_mode

Example 4: convert T5 model with beam search in two steps:
    python -m models.t5.convert_to_onnx -m t5-small
    python convert_generation.py -m t5-small --model_type t5             \
        --decoder_onnx ./onnx_models/t5-small_decoder.onnx               \
        --encoder_decoder_init_onnx ./onnx_models/t5-small_encoder.onnx  \
        --output ./onnx_models/t5_small_beam_search.onnx

Example 5: convert T5 model with beam search. All in one step:
    python convert_generation.py -m t5-small --model_type t5 --output t5_small_beam_search.onnx

Example 6: convert T5 model with beam search containing specific cuda optimizations. All in one step:
    python convert_generation.py -m t5-small --model_type t5 --output t5_small_beam_search.onnx   \
        --use_gpu --past_present_share_buffer --use_decoder_masked_attention

Example 7: convert MT5 model with external data file like mt5-base-beamsearch.onnx.data in below example.
    python convert_generation.py -m google/mt5-base --model_type mt5 --output mt5-base-beamsearch.onnx -e

Example 8: convert gpt2 model with greedy search:
    python convert_generation.py -m gpt2 --output gpt2_greedy_search.onnx --num_beams 1 --num_return_sequences 1

Example 9: convert gpt2 model with sampling:
    python convert_generation.py -m gpt2 --output gpt2_sampling.onnx --num_beams 1 --num_return_sequences 1 --top_p 0.6
"""

import argparse
import logging
import math
import os
import time
from enum import Enum
from pathlib import Path
from typing import Any

import numpy as np
import onnx
import torch
from benchmark_helper import Precision, setup_logger
from fusion_utils import NumpyHelper
from onnx import GraphProto, ModelProto, TensorProto
from onnx_model import OnnxModel
from transformers import (
    GPT2Config,
    GPT2LMHeadModel,
    GPT2Tokenizer,
    MT5Config,
    MT5ForConditionalGeneration,
    T5Config,
    T5ForConditionalGeneration,
    T5Tokenizer,
)

from onnxruntime import (
    GraphOptimizationLevel,
    InferenceSession,
    SessionOptions,
    get_available_providers,
)
from onnxruntime.transformers.models.gpt2.convert_to_onnx import (
    main as convert_gpt2_to_onnx,
)
from onnxruntime.transformers.models.gpt2.gpt2_helper import PRETRAINED_GPT2_MODELS
from onnxruntime.transformers.models.t5.convert_to_onnx import (
    export_onnx_models as export_t5_onnx_models,
)
from onnxruntime.transformers.models.t5.t5_helper import (
    PRETRAINED_MT5_MODELS,
    PRETRAINED_T5_MODELS,
)

logger = logging.getLogger("")


class GenerationType(Enum):
    BEAMSEARCH = "beam_search"
    GREEDYSEARCH = "greedy_search"
    SAMPLING = "sampling"

    def __str__(self):
        return self.value


def parse_arguments(argv: list[str] | None = None) -> argparse.Namespace:
    """Parse arguments

    Args:
        argv (Optional[List[str]], optional): _description_. Defaults to None.

    Returns:
        argparse.Namespace: Parsed arguments.
    """
    parser = argparse.ArgumentParser()

    input_group = parser.add_argument_group("Input options")

    input_group.add_argument(
        "-m",
        "--model_name_or_path",
        required=True,
        type=str,
        help="Pytorch model checkpoint path, or pretrained model name in the list: "
        + ", ".join(PRETRAINED_GPT2_MODELS + PRETRAINED_T5_MODELS + PRETRAINED_MT5_MODELS),
    )

    input_group.add_argument(
        "--model_type",
        required=False,
        type=str,
        default="gpt2",
        choices=["gpt2", "t5", "mt5"],
        help="Model type (default is gpt2) in the list: " + ", ".join(["gpt2", "t5", "mt5"]),
    )

    input_group.add_argument(
        "--cache_dir",
        required=False,
        type=str,
        default=os.path.join(".", "cache_models"),
        help="Directory to cache pre-trained models",
    )

    input_group.add_argument(
        "--decoder_onnx",
        required=False,
        type=str,
        default="",
        help="Path of onnx model for decoder. Specify it when you have exported the model.",
    )

    input_group.add_argument(
        "--encoder_decoder_init_onnx",
        required=False,
        type=str,
        default="",
        help="Path of ONNX model for encoder and decoder initialization. Specify it when you have exported the model.",
    )

    parser.add_argument(
        "--verbose",
        required=False,
        action="store_true",
        help="Print more information",
    )
    parser.set_defaults(verbose=False)

    output_group = parser.add_argument_group("Output options")

    output_group.add_argument(
        "--output",
        required=True,
        type=str,
        help="Output path for onnx model with beam search.",
    )

    output_group.add_argument(
        "-p",
        "--precision",
        required=False,
        type=str,
        default=Precision.FLOAT32.value,
        choices=[Precision.FLOAT32.value, Precision.FLOAT16.value],
        help="Precision of model to run. fp32 for full precision, fp16 for half or mixed precision",
    )

    output_group.add_argument(
        "-b",
        "--op_block_list",
        required=False,
        nargs="*",
        default=["auto"],
        help="Disable certain onnx operators when exporting model to onnx format. When using default"
        'value for gpt2 type of model fp16 precision, it will be set to ["Add", "LayerNormalization",'
        ' "SkipLayerNormalization", "FastGelu"]. Other situation, it will be set to []',
    )

    output_group.add_argument(
        "-e",
        "--use_external_data_format",
        required=False,
        action="store_true",
        help="save external data for model > 2G",
    )
    output_group.set_defaults(use_external_data_format=False)

    output_group.add_argument(
        "-s",
        "--run_shape_inference",
        required=False,
        action="store_true",
        help="run shape inference",
    )
    output_group.set_defaults(run_shape_inference=False)

    output_group.add_argument(
        "-dpvs",
        "--disable_pad_vocab_size",
        required=False,
        action="store_true",
        help="Do not pad logits MatMul weight to be a multiple of 8 along the dimension where dim value is"
        " the vocab size. The logits MatMul may hence be of poor performance for fp16 precision.",
    )
    output_group.set_defaults(disable_pad_vocab_size=False)

    output_group.add_argument(
        "-dsgd",
        "--disable_separate_gpt2_decoder_for_init_run",
        required=False,
        action="store_true",
        help="Do not create separate decoder subgraphs for initial and remaining runs. This does not allow "
        "for optimizations based on sequence lengths in each subgraph",
    )
    output_group.set_defaults(disable_separate_gpt2_decoder_for_init_run=False)

    output_group.add_argument(
        "-i",
        "--disable_shared_initializers",
        required=False,
        action="store_true",
        help="do not share initializers in encoder and decoder for T5 or in the init decoder and decoder for "
        "GPT2. It will increase memory usage of t5/mt5/gpt2 models.",
    )
    output_group.set_defaults(disable_shared_initializers=False)

    output_group.add_argument(
        "--encoder_decoder_init",
        required=False,
        action="store_true",
        help="Add decoder initialization to encoder for T5 model. This is legacy format that will be deprecated.",
    )
    output_group.set_defaults(encoder_decoder_init=False)

    model_group = parser.add_argument_group("Beam search parameters that stored in the output model")

    model_group.add_argument(
        "--output_sequences_scores",
        required=False,
        action="store_true",
        help="output sequences scores",
    )
    model_group.set_defaults(output_sequences_scores=False)

    model_group.add_argument(
        "--output_token_scores",
        required=False,
        action="store_true",
        help="output token scores",
    )
    model_group.set_defaults(output_token_scores=False)

    model_group.add_argument("--early_stopping", required=False, action="store_true")
    model_group.set_defaults(early_stopping=False)

    model_group.add_argument(
        "--no_repeat_ngram_size",
        type=int,
        required=False,
        default=0,
        help="No repeat ngram size",
    )

    model_group.add_argument(
        "--vocab_mask",
        required=False,
        action="store_true",
        help="Enable vocab_mask. This mask applies only to every generated token to filter some bad words.",
    )
    model_group.set_defaults(vocab_mask=False)

    model_group.add_argument(
        "--past_present_share_buffer",
        required=False,
        action="store_true",
        help="Use shared buffer for past and present, currently work for gpt2 greedy/sampling search.",
    )
    model_group.set_defaults(past_present_share_buffer=False)

    model_group.add_argument(
        "--use_decoder_masked_attention",
        required=False,
        action="store_true",
        help="Uses `DecoderMaskedSelfAttention` or `DecoderMaskedMultiHeadAttention` to optimize the decoding Attention computation. "
        "Must be used with `past_present_share_buffer`. Currently, only Attention head sizes of 32, 64 and 128 are supported.",
    )
    model_group.set_defaults(use_decoder_masked_attention=False)

    model_group.add_argument(
        "--prefix_vocab_mask",
        required=False,
        action="store_true",
        help="Enable prefix_vocab_mask. This mask can be used to filter bad words in the first generated token only",
    )
    model_group.set_defaults(prefix_vocab_mask=False)

    model_group.add_argument(
        "--custom_attention_mask",
        required=False,
        action="store_true",
        help="Enable custom_attention_mask. This mask can be used to replace default encoder attention mask",
    )
    model_group.set_defaults(custom_attention_mask=False)

    model_group.add_argument(
        "--presence_mask",
        required=False,
        action="store_true",
        help="Presence mask for custom sampling",
    )
    model_group.set_defaults(presence_mask=False)

    model_group.add_argument(
        "--seed",
        required=False,
        action="store_true",
        help="Random seed for sampling op",
    )
    model_group.set_defaults(seed=False)

    beam_parameters_group = parser.add_argument_group(
        "Beam search parameters not stored in the output model, for testing parity and performance"
    )

    beam_parameters_group.add_argument("--min_length", type=int, required=False, default=1, help="Min sequence length")

    beam_parameters_group.add_argument("--max_length", type=int, required=False, default=50, help="Max sequence length")

    beam_parameters_group.add_argument("--num_beams", type=int, required=False, default=4, help="Beam size")

    beam_parameters_group.add_argument(
        "--num_return_sequences",
        type=int,
        required=False,
        default=1,
        help="Number of return sequence <= num_beams",
    )

    beam_parameters_group.add_argument(
        "--length_penalty",
        type=float,
        required=False,
        default=1,
        help="Positive. >1 to penalize and <1 to encourage short sentence.",
    )

    beam_parameters_group.add_argument(
        "--repetition_penalty",
        type=float,
        required=False,
        default=1,
        help="Positive. >1 to penalize and <1 to encourage.",
    )

    beam_parameters_group.add_argument(
        "--temperature",
        type=float,
        required=False,
        default=1.0,
        help="The value used to module the next token probabilities.",
    )

    beam_parameters_group.add_argument(
        "--top_p",
        type=float,
        required=False,
        default=1.0,
        help="Top P for sampling",
    )

    beam_parameters_group.add_argument(
        "--filter_value",
        type=float,
        required=False,
        default=-float("Inf"),
        help="Filter value for Top P sampling",
    )

    beam_parameters_group.add_argument(
        "--min_tokens_to_keep",
        type=int,
        required=False,
        default=1,
        help="Minimum number of tokens we keep per batch example in the output.",
    )

    beam_parameters_group.add_argument(
        "--presence_penalty",
        type=float,
        required=False,
        default=0.0,
        help="presence penalty for custom sampling.",
    )

    beam_parameters_group.add_argument(
        "--custom",
        type=int,
        required=False,
        default=0,
        help="If 1 customized top P logic is applied",
    )

    beam_parameters_group.add_argument(
        "--vocab_size",
        type=int,
        required=False,
        default=-1,
        help="Vocab_size of the underlying model used to decide the shape of vocab mask",
    )

    beam_parameters_group.add_argument(
        "--eos_token_id",
        type=int,
        required=False,
        default=-1,
        help="custom eos_token_id for generating model with existing onnx encoder/decoder",
    )

    beam_parameters_group.add_argument(
        "--pad_token_id",
        type=int,
        required=False,
        default=-1,
        help="custom pad_token_id for generating model with existing onnx encoder/decoder",
    )

    test_group = parser.add_argument_group("Other options for testing parity and performance")

    test_group.add_argument(
        "--use_sln_strict_mode",
        required=False,
        action="store_true",
        help="Enable strict mode for SLN in CUDA provider. This ensures a better accuracy but will be slower.",
    )
    test_group.set_defaults(use_sln_strict_mode=False)

    test_group.add_argument(
        "--use_gpu",
        required=False,
        action="store_true",
        help="use GPU for inference. Required for fp16.",
    )
    test_group.set_defaults(use_gpu=False)

    test_group.add_argument(
        "--disable_parity",
        required=False,
        action="store_true",
        help="do not run parity test",
    )
    test_group.set_defaults(disable_parity=False)

    test_group.add_argument(
        "--disable_perf_test",
        required=False,
        action="store_true",
        help="do not run perf test",
    )
    test_group.set_defaults(disable_perf_test=False)

    test_group.add_argument(
        "--torch_performance",
        required=False,
        action="store_true",
        help="test PyTorch performance",
    )
    test_group.set_defaults(torch_performance=False)

    test_group.add_argument(
        "--total_runs",
        required=False,
        type=int,
        default=1,
        help="Number of times of inference for latency measurement",
    )

    test_group.add_argument(
        "--save_test_data",
        required=False,
        action="store_true",
        help="save test data for onnxruntime_perf_test tool",
    )
    test_group.set_defaults(save_test_data=False)

    args = parser.parse_args(argv)

    return args


def gpt2_to_onnx(args: argparse.Namespace):
    """Convert GPT-2 model to onnx

    Args:
        args (argparse.Namespace): arguments parsed from command line
    """
    model_name = args.model_name_or_path

    arguments = [
        "--model_name_or_path",
        model_name,
        "--output",
        args.decoder_onnx,
        "--optimize_onnx",
        "--precision",
        args.precision,
        "--test_runs",
        "1",
        "--test_cases",
        "10",
        "--overwrite",  # Overwrite onnx file if existed
    ]
    if args.cache_dir:
        arguments.extend(["--cache_dir", args.cache_dir])
    if args.use_gpu:
        arguments.append("--use_gpu")
    if args.use_external_data_format:
        arguments.append("--use_external_data_format")

    if len(args.op_block_list):
        arguments.extend(["--op_block_list"])
        arguments.extend(args.op_block_list)

    if args.precision == Precision.FLOAT16.value:
        assert args.use_gpu, "fp16 or mixed precision model cannot run in CPU. Please add --use_gpu"
        # TODO(tianleiwu): Use auto mixed precision for fp16 conversion: arguments.append('--auto_mixed_precision')
        #       Need change cuda kernel to support a combination of fp32 logits and fp16 past state.
        #       Currently logits and past state shall be same data type.

    if args.verbose:
        logger.info(f"arguments for convert_to_onnx:{arguments}")

    convert_gpt2_to_onnx(argv=arguments)


def t5_to_onnx(args: argparse.Namespace):
    """Convert T5 model to onnx

    Args:
        args (argparse.Namespace): arguments parsed from command line
    """
    paths = export_t5_onnx_models(
        model_name_or_path=args.model_name_or_path,
        cache_dir=args.cache_dir,
        output_dir=Path(args.output).parent,
        use_gpu=args.use_gpu,
        use_external_data_format=args.use_external_data_format,
        optimize_onnx=(args.precision != Precision.FLOAT16.value),
        precision=args.precision,
        verbose=False,
        use_decoder_start_token=False,
        overwrite=True,
        disable_auto_mixed_precision=False,
        use_int32_inputs=True,
        model_type=args.model_type,
        encoder_decoder_init=args.encoder_decoder_init,
        force_fp16_io=(args.precision == Precision.FLOAT16.value),  # required by BeamSearch op implementation.
    )

    logger.debug(f"onnx model for encoder: {paths[0]}")
    logger.debug(f"onnx model for decoder: {paths[1]}")
    args.encoder_decoder_init_onnx = paths[0]
    args.decoder_onnx = paths[1]


def shape_inference(onnx_path: str, use_external_data_format: bool = True):
    """Shape inference on an onnx file, which will be overwritten.

    Args:
        onnx_path (str): Path of onnx model
        use_external_data_format(bool): output tensors to external data or not.
    """
    # Run symbolic shape inference to walk around ORT shape inference issue for subgraph.
    from onnxruntime.tools.symbolic_shape_infer import SymbolicShapeInference  # noqa: PLC0415

    model = onnx.load_model(onnx_path, load_external_data=True)
    out = SymbolicShapeInference.infer_shapes(model, auto_merge=True, guess_output_rank=False)
    if out:
        OnnxModel.save(out, onnx_path, save_as_external_data=use_external_data_format)
    else:
        logger.warning("Failed to run symbolic shape inference on the model.")


def pad_weights_of_logits_matmul(onnx_path: str, use_external_data_format: bool = True) -> bool:
    """Pad the logits MatMul weight in the provided decoder model, which will be overwritten.

    Args:
        onnx_path (str): Path of onnx model
        use_external_data_format(bool): output tensors to external data or not.
    """
    decoder_model_proto = onnx.load_model(onnx_path, load_external_data=True)

    logits_output_name = decoder_model_proto.graph.output[0].name

    decoder_model = OnnxModel(decoder_model_proto)

    output_name_to_node = decoder_model.output_name_to_node()
    assert logits_output_name in output_name_to_node

    matmul_node = output_name_to_node[logits_output_name]
    # Sanity check - the logits need to be produced by a MatMul node
    if matmul_node.op_type != "MatMul":
        return False

    # The logits MatMul weight MUST be an initializer (or)
    # it MUST be flowing through a Transpose whose input is
    # an initializer
    pad_along_axis_1 = True
    logits_weight = decoder_model.get_initializer(matmul_node.input[1])
    if logits_weight is None:
        transpose_before_matmul = decoder_model.match_parent(matmul_node, "Transpose", 1)

        if transpose_before_matmul is None:
            return False

        logits_weight = decoder_model.get_initializer(transpose_before_matmul.input[0])

        if logits_weight is None:
            return False

        pad_along_axis_1 = False

    # The logits MatMul weight MUST be fp16
    if logits_weight.data_type != TensorProto.DataType.FLOAT16:
        return False

    # The logits MatMul weight MUST be 2-dimensional
    if len(logits_weight.dims) != 2:
        return False

    # Pad and over-write the initializer (if needed)
    actual_vocab_size = logits_weight.dims[1]

    if (actual_vocab_size % 8) == 0:
        # Already "padded"
        return True

    padded_vocab_size = math.ceil(actual_vocab_size / 8) * 8
    padding = padded_vocab_size - actual_vocab_size

    # TODO(hasesh): Handle cases where the fp16 data is stored in the
    # non-raw data field
    if logits_weight.raw_data:
        if pad_along_axis_1:
            padding_data = np.zeros((logits_weight.dims[0], padding), dtype=np.float16)
            weight_with_padding = np.concatenate((NumpyHelper.to_array(logits_weight), padding_data), axis=1)
            logits_weight.dims[1] = padded_vocab_size
        else:
            padding_data = np.zeros((padding, logits_weight.dims[1]), dtype=np.float16)
            weight_with_padding = np.concatenate((NumpyHelper.to_array(logits_weight), padding_data), axis=0)
            logits_weight.dims[0] = padded_vocab_size

        logits_weight.raw_data = weight_with_padding.tobytes()
    else:
        return False

    # Save the model
    OnnxModel.save(decoder_model_proto, onnx_path, save_as_external_data=use_external_data_format)
    return True


def create_ort_session(model_path: str, use_gpu: bool, use_sln_strict_mode: bool) -> InferenceSession:
    """Create OnnxRuntime session.

    Args:
        model_path (str): onnx model path
        use_gpu (bool): use GPU or not
        use_sln_strict_mode (bool): use strict mode for skip layer normalization or not

    Raises:
        RuntimeError: CUDAExecutionProvider is not available when --use_gpu is specified.

    Returns:
        onnxruntime.InferenceSession: The created session.
    """
    sess_options = SessionOptions()
    sess_options.graph_optimization_level = GraphOptimizationLevel.ORT_DISABLE_ALL
    execution_providers = ["CUDAExecutionProvider", "CPUExecutionProvider"] if use_gpu else ["CPUExecutionProvider"]
    if use_gpu:
        if "CUDAExecutionProvider" not in get_available_providers():
            raise RuntimeError("CUDAExecutionProvider is not available for --use_gpu!")
        else:
            logger.info("use CUDAExecutionProvider")
        if use_sln_strict_mode:
            cuda_provider_options = {"enable_skip_layer_norm_strict_mode": True}
            provider_options = {"CUDAExecutionProvider": cuda_provider_options}
            execution_providers = [
                (name, provider_options[name]) if name in provider_options else name for name in execution_providers
            ]

    ort_session = InferenceSession(model_path, sess_options, providers=execution_providers)
    return ort_session


def verify_gpt2_subgraph(graph: onnx.GraphProto, precision: Precision):
    """Verify GPT-2 subgraph

    Args:
        graph (onnx.GraphProto): onnx graph of GPT-2
        precision (Precision): Precision (FLOAT16 or FLOAT32) of the model.

    Raises:
        ValueError: Number of inputs not expected.
        ValueError: Input name is not expected.
        ValueError: Input data type is not expected.
        ValueError: Number of outputs not expected.
        ValueError: Output name is not expected.
        ValueError: Output data type is not expected.
    """
    is_float16 = precision == Precision.FLOAT16.value

    input_count = len(graph.input)
    layer_count = input_count - 3
    assert layer_count >= 1

    expected_inputs = ["input_ids", "position_ids", "attention_mask"] + [f"past_{i}" for i in range(layer_count)]
    if len(graph.input) != len(expected_inputs):
        raise ValueError(f"Number of inputs expected to be {len(expected_inputs)}. Got {len(graph.input)}")

    for i, expected_input in enumerate(expected_inputs):
        if graph.input[i].name != expected_input:
            raise ValueError(f"Input {i} is expected to be {expected_input}. Got {graph.input[i].name}")

        expected_type = TensorProto.INT32
        if i >= 3:
            expected_type = TensorProto.FLOAT16 if is_float16 else TensorProto.FLOAT

        input_type = graph.input[i].type.tensor_type.elem_type
        if input_type != expected_type:
            raise ValueError(f"Input {i} is expected to have onnx data type {expected_type}. Got {input_type}")
    logger.info("Verifying GPT-2 graph inputs: name and data type are good.")

    expected_outputs = ["logits"] + [f"present_{i}" for i in range(layer_count)]
    if len(graph.output) != len(expected_outputs):
        raise ValueError(f"Number of outputs expected to be {len(expected_outputs)}. Got {len(graph.output)}")

    for i, expected_output in enumerate(expected_outputs):
        if graph.output[i].name != expected_output:
            raise ValueError(f"Output {i} is expected to be {expected_output}. Got {graph.output[i].name}")

        expected_type = TensorProto.FLOAT16 if is_float16 else TensorProto.FLOAT
        output_type = graph.output[i].type.tensor_type.elem_type
        if output_type != expected_type:
            raise ValueError(f"Input {i} is expected to have onnx data type {expected_type}. Got {output_type}")
    logger.info("Verifying GPT-2 graph outputs: name and data type are good.")

    # TODO(tianleiwu): verify shapes of inputs and outputs.
    return


def verify_t5_decoder_subgraph(graph: onnx.GraphProto, precision: Precision):
    """Verify T5 decoder subgraph

    Args:
        graph (onnx.GraphProto): onnx graph of T5 decoder
        precision (Precision): Precision (FLOAT16 or FLOAT32) of the model.

    Raises:
        ValueError: Number of inputs not expected.
        ValueError: Input name is not expected.
        ValueError: Input data type is not expected.
        ValueError: Number of outputs not expected.
        ValueError: Output name is not expected.
        ValueError: Output data type is not expected.
    """
    is_float16 = precision == Precision.FLOAT16.value
    float_type = TensorProto.FLOAT16 if is_float16 else TensorProto.FLOAT

    input_count = len(graph.input)
    layer_count = (input_count - 2) // 4
    assert layer_count >= 1

    # Expect inputs:
    #   input_ids: int32 (B, 1)
    #   encoder_attention_mask: int32 (B, encode_sequence_length)

    #   past_key_self_0: (B, num_heads, past_decode_sequence_length, head_size)
    #   past_value_self_0: (B, num_heads, past_decode_sequence_length, head_size)
    #   ... (for each self attention layer)

    #   past_key_cross_0: (B, num_heads, encode_sequence_length, head_size)
    #   past_value_cross_0: (B, num_heads, encode_sequence_length, head_size)
    #   ... (for each cross attention layer)

    # TODO: encoder_hidden_states is optional
    expected_inputs = ["input_ids", "encoder_attention_mask"]
    for i in range(layer_count):
        expected_inputs.append(f"past_key_self_{i}")
        expected_inputs.append(f"past_value_self_{i}")
    for i in range(layer_count):
        expected_inputs.append(f"past_key_cross_{i}")
        expected_inputs.append(f"past_value_cross_{i}")

    if len(graph.input) != len(expected_inputs):
        raise ValueError(f"Number of inputs expected to be {len(expected_inputs)}. Got {len(graph.input)}")

    for i, expected_input in enumerate(expected_inputs):
        if graph.input[i].name != expected_input:
            raise ValueError(f"Input {i} is expected to be {expected_input}. Got {graph.input[i].name}")

        expected_type = TensorProto.INT32 if i < 2 else float_type
        input_type = graph.input[i].type.tensor_type.elem_type
        if input_type != expected_type:
            raise ValueError(f"Input {i} is expected to have onnx data type {expected_type}. Got {input_type}")

    # Expect outputs:
    #   logits:               (B, 1, vocab_size)
    #   present_key_self_0:   (B, num_heads, past_decode_sequence_length + 1, head_size)
    #   present_value_self_0: (B, num_heads, past_decode_sequence_length + 1, head_size)
    #                     ... (for each self attention layer)
    expected_outputs = ["logits"]
    for i in range(layer_count):
        expected_outputs.append(f"present_key_self_{i}")
        expected_outputs.append(f"present_value_self_{i}")

    if len(graph.output) != len(expected_outputs):
        raise ValueError(f"Number of outputs expected to be {len(expected_outputs)}. Got {len(graph.output)}")

    for i, expected_output in enumerate(expected_outputs):
        if graph.output[i].name != expected_output:
            raise ValueError(f"Output {i} is expected to be {expected_output}. Got {graph.output[i].name}")
        output_type = graph.output[i].type.tensor_type.elem_type
        if output_type != float_type:
            raise ValueError(f"Output {i} is expected to have onnx data type {float_type}. Got {output_type}")


def verify_t5_encoder_decoder_init_subgraph(graph: onnx.GraphProto, precision: Precision):
    """Verify T5 decoder subgraph

    Args:
        graph (onnx.GraphProto): onnx graph of T5 decoder
        precision (Precision): Precision (FLOAT16 or FLOAT32) of the model.

    Raises:
        ValueError: Number of inputs not expected.
        ValueError: Input name is not expected.
        ValueError: Input data type is not expected.
        ValueError: Number of outputs not expected.
        ValueError: Output name is not expected.
        ValueError: Output data type is not expected.
    """
    is_float16 = precision == Precision.FLOAT16.value
    new_format = "cross" in graph.output[0].name

    # Expect 3 inputs:
    #   encoder_input_ids:      int32 (B, encode_sequence_length)
    #   encoder_attention_mask: int32 (B, encode_sequence_length)
    #   decoder_input_ids:      int32 (B, 1)
    expected_inputs = [
        "encoder_input_ids",
        "encoder_attention_mask",
        "decoder_input_ids",
    ]
    if new_format:
        expected_inputs = expected_inputs[:2]
    if len(graph.input) != len(expected_inputs):
        raise ValueError(f"Number of inputs expected to be {len(expected_inputs)}. Got {len(graph.input)}")

    for i, expected_input in enumerate(expected_inputs):
        if graph.input[i].name != expected_input:
            raise ValueError(f"Input {i} is expected to be {expected_input}. Got {graph.input[i].name}")

        expected_type = TensorProto.INT32
        input_type = graph.input[i].type.tensor_type.elem_type
        if input_type != expected_type:
            raise ValueError(f"Input {i} is expected to have onnx data type {expected_type}. Got {input_type}")

    if new_format:
        assert len(graph.output) % 2 == 0
        layer_count = len(graph.output) // 2
        assert layer_count >= 1

        # Expected outputs:
        #   present_key_cross_0:   (B, num_heads, encode_sequence_length, head_size)
        #   present_value_cross_0: (B, num_heads, encode_sequence_length, head_size)
        #                      ... (for each cross attention layer)
        expected_outputs = []
        for i in range(layer_count):
            expected_outputs.append(f"present_key_cross_{i}")
            expected_outputs.append(f"present_value_cross_{i}")
    else:
        logger.warning("This format is deprecated. Please export T5 encoder in new format with only cross outputs.")
        assert (len(graph.output) - 2) % 4 == 0
        layer_count = (len(graph.output) - 2) // 4
        assert layer_count >= 1

        # Expected outputs:
        #   logits:                (B, 1, vocab_size)
        #   encoder_hidden_states: (B, encode_sequence_length, encoder_hidden_size)
        #   present_key_self_0:    (B, num_heads, 1, head_size)
        #   present_value_self_0:  (B, num_heads, 1, head_size)
        #                      ... (for each self attention layer)
        #   present_key_cross_0:   (B, num_heads, encode_sequence_length, head_size)
        #   present_value_cross_0: (B, num_heads, encode_sequence_length, head_size)
        #                      ... (for each cross attention layer)
        expected_outputs = ["logits", "encoder_hidden_states"]
        for i in range(layer_count):
            expected_outputs.append(f"present_key_self_{i}")
            expected_outputs.append(f"present_value_self_{i}")
        for i in range(layer_count):
            expected_outputs.append(f"present_key_cross_{i}")
            expected_outputs.append(f"present_value_cross_{i}")

    if len(graph.output) != len(expected_outputs):
        raise ValueError(f"Number of outputs expected to be {len(expected_outputs)}. Got {len(graph.output)}")

    for i, expected_output in enumerate(expected_outputs):
        if graph.output[i].name != expected_output:
            raise ValueError(f"Output {i} is expected to be {expected_output}. Got {graph.output[i].name}")

        expected_type = TensorProto.FLOAT16 if is_float16 else TensorProto.FLOAT
        output_type = graph.output[i].type.tensor_type.elem_type
        if output_type != expected_type:
            raise ValueError(f"Output {i} is expected to have onnx data type {expected_type}. Got {output_type}")

    logger.info("T5 encoder graph verified: name and data type of inputs and outputs are good.")


def remove_shared_initializers(
    graph1: GraphProto,
    graph2: GraphProto,
    shared_prefix: str = "shared_",
    min_elements: int = 1024,
    signature_cache1: dict | None = None,
    signature_cache2: dict | None = None,
):
    """Remove initializers with same value from two graphs.

    Args:
        graph1 (GraphProto): the first graph to process
        graph2 (GraphProto): the second graph to process
        shared_prefix (str): add prefix to the shared initializers among two graphs
        min_elements (int, optional): minimal number of elements for initializers to be considered. Defaults to 1024.
        signature_cache1 (dict): Optional dictionary to store data signatures of tensors in graph1 in order to speed up comparison
        signature_cache2 (dict): Optional dictionary to store data signatures of tensors in graph2 in order to speed up comparison
    """

    mapping_initializers_1 = {}
    mapping_initializers_2 = {}
    shared_initializers_1 = []
    shared_initializers_2 = []
    shared_initializers_names = []

    for initializer1 in graph1.initializer:
        if not (initializer1.dims and sum(initializer1.dims) >= min_elements):
            continue

        for initializer2 in graph2.initializer:
            if not (initializer2.dims and sum(initializer2.dims) >= min_elements):
                continue

            if OnnxModel.has_same_value(initializer1, initializer2, signature_cache1, signature_cache2):
                mapping_initializers_1[initializer1.name] = shared_prefix + initializer2.name
                shared_initializers_1.append(initializer1)

                if initializer2.name not in mapping_initializers_2:
                    shared_name = shared_prefix + initializer2.name
                    mapping_initializers_2[initializer2.name] = shared_name
                    shared_initializers_2.append(initializer2)
                    shared_initializers_names.append(shared_name)
                break

    logger.debug(f"shared initializers:{shared_initializers_names}")

    # Make sure new name does not exist in graph 1
    for node in graph1.node:
        for j in range(len(node.input)):
            if node.input[j] in shared_initializers_names:
                raise RuntimeError(f"name is found in graph 1: {node.input[j]}")

    # Make sure new name does not exist in graph 2
    for node in graph2.node:
        for j in range(len(node.input)):
            if node.input[j] in shared_initializers_names:
                raise RuntimeError(f"name is found in graph 2: {node.input[j]}")

    # Remove shared initializers from graph 2
    for initializer in shared_initializers_2:
        graph2.initializer.remove(initializer)

    # Rename value info for old names in graph 2
    for value_info in graph2.value_info:
        if value_info.name in mapping_initializers_2:
            value_info.name = mapping_initializers_2[value_info.name]

    # Rename nodes inputs in graph 2:
    for node in graph2.node:
        for j in range(len(node.input)):
            if node.input[j] in mapping_initializers_2:
                new_name = mapping_initializers_2[node.input[j]]
                logger.debug(f"graph 2 rename node {node.name} input {j} from {node.input[j]} to {new_name}")
                node.input[j] = new_name

    #  Remove shared initializers from graph 1
    for initializer in shared_initializers_1:
        graph1.initializer.remove(initializer)

    # Rename value info for old names in graph 1
    for value_info in graph1.value_info:
        if value_info.name in mapping_initializers_1:
            value_info.name = mapping_initializers_1[value_info.name]

    # Rename nodes inputs in graph 1:
    for node in graph1.node:
        for j in range(len(node.input)):
            if node.input[j] in mapping_initializers_1:
                new_name = mapping_initializers_1[node.input[j]]
                logger.debug(f"graph 1 rename node {node.name} input {j} from {node.input[j]} to {new_name}")
                node.input[j] = new_name

    # Rename shared initializers in graph 2
    for initializer in shared_initializers_2:
        initializer.name = mapping_initializers_2[initializer.name]

    for initializer in shared_initializers_2:
        shape = onnx.numpy_helper.to_array(initializer).shape
        value_info = onnx.helper.make_tensor_value_info(initializer.name, initializer.data_type, shape)
        # Need add value_info for initializers moved to parent graph. Otherwise, ORT will fail.
        graph1.value_info.append(value_info)
        graph2.value_info.append(value_info)

    return shared_initializers_2


def get_shared_initializers(encoder_model: ModelProto, decoder_model: ModelProto):
    encoder = OnnxModel(encoder_model)
    decoder = OnnxModel(decoder_model)
    encoder.add_prefix_to_names("e_")
    decoder.add_prefix_to_names("d_")
    signature_cache1, signature_cache2 = {}, {}
    encoder.remove_duplicated_initializer(signature_cache1)
    decoder.remove_duplicated_initializer(signature_cache2)
    initializers = remove_shared_initializers(
        decoder.model.graph,
        encoder.model.graph,
        shared_prefix="s_",
        signature_cache1=signature_cache1,
        signature_cache2=signature_cache2,
    )
    return initializers


def move_initializers(
    graph: GraphProto,
    min_elements: int = 1024,
) -> list[TensorProto]:
    """Remove initializers of a graph, when they have number of elements larger than a threshold.

    Args:
        graph (GraphProto): the graph.
        min_elements (int, optional): minimal number of elements for initializers to be considered. Defaults to 1024.

    Returns:
        List[TensorProto]: initializers that are removed from the graph.
    """
    moved_initializers = []
    for tensor in graph.initializer:
        if not (tensor.dims and sum(tensor.dims) >= min_elements):
            continue
        moved_initializers.append(tensor)

    for initializer in moved_initializers:
        graph.initializer.remove(initializer)

    # Add type info, otherwise ORT will raise error: "input arg (*) does not have type information set by parent node."
    for initializer in moved_initializers:
        shape = onnx.numpy_helper.to_array(initializer).shape
        value_info = onnx.helper.make_tensor_value_info(initializer.name, initializer.data_type, shape)
        graph.value_info.append(value_info)

    return moved_initializers


def _attribute_to_pair(attribute):
    """
    Convert attribute to kwarg format for use with onnx.helper.make_node.
        :parameter attribute: attribute in AttributeProto format.
        :return: attribute in {key: value} format.
    """
    if attribute.type == 0:
        raise ValueError(f"attribute {attribute.name} does not have type specified.")

    # Based on attribute type definitions from AttributeProto
    # definition in https://github.com/onnx/onnx/blob/master/onnx/onnx.proto
    if attribute.type == 1:
        value = attribute.f
    elif attribute.type == 2:
        value = attribute.i
    elif attribute.type == 3:
        value = attribute.s
    elif attribute.type == 4:
        value = attribute.t
    elif attribute.type == 5:
        value = attribute.g
    elif attribute.type == 6:
        value = attribute.floats
    elif attribute.type == 7:
        value = attribute.ints
    elif attribute.type == 8:
        value = attribute.strings
    elif attribute.type == 9:
        value = attribute.tensors
    elif attribute.type == 10:
        value = attribute.graphs
    else:
        raise ValueError(f"attribute {attribute.name} has unsupported type {attribute.type}.")

    return (attribute.name, value)


def kwargs_of(node):
    kwargs = {}
    for attr in node.attribute:
        (key, value) = _attribute_to_pair(attr)
        kwargs.update({key: value})
    if node.domain:
        kwargs.update({"domain": node.domain})
    return kwargs


def shape_of(vi):
    return tuple([d.dim_param if (d.dim_param) else d.dim_value for d in vi.type.tensor_type.shape.dim])


def update_decoder_subgraph_past_present_share_buffer(subg: GraphProto):
    input_past_0 = 3
    output_past_0 = 1
    new_inputs = []
    for i, vi in enumerate(subg.input):
        if i >= input_past_0:
            shape = shape_of(vi)
            vi = onnx.helper.make_tensor_value_info(  # noqa: PLW2901
                vi.name,
                elem_type=vi.type.tensor_type.elem_type,
                shape=[shape[0], shape[1], shape[2], "max_seq_len", shape[4]],
            )
        new_inputs.extend([vi])
    new_inputs.extend([onnx.helper.make_tensor_value_info("past_sequence_length", onnx.TensorProto.INT32, shape=[1])])
    subg.ClearField("input")
    subg.input.extend(new_inputs)

    new_outputs = []
    for i, vi in enumerate(subg.output):
        if i >= output_past_0:
            shape = shape_of(vi)
            vi = onnx.helper.make_tensor_value_info(  # noqa: PLW2901
                vi.name,
                elem_type=vi.type.tensor_type.elem_type,
                shape=[shape[0], shape[1], shape[2], "max_seq_len", shape[4]],
            )
        new_outputs.extend([vi])
    subg.ClearField("output")
    subg.output.extend(new_outputs)

    new_nodes = []
    for node in subg.node:
        new_node = node
        if node.op_type == "Attention":
            kwargs = kwargs_of(node)
            kwargs.update({"past_present_share_buffer": 1})
            nis = []
            nis.extend(node.input)
            while len(nis) < 6:
                nis.extend([""])
            if len(nis) < 7:
                nis.extend(["past_sequence_length"])
            new_node = onnx.helper.make_node("Attention", nis, node.output, name=node.name, **kwargs)
        new_nodes.extend([new_node])
    subg.ClearField("node")
    subg.node.extend(new_nodes)
    return subg


def update_decoder_subgraph_use_decoder_masked_attention(
    subg: GraphProto, is_beam_search: bool, switch_attention: bool
) -> bool:
    """Update the Attention nodes to DecoderMaskedSelfAttention.

    Args:
        subg (GraphProto): GraphProto of the decoder subgraph
        is_beam_search (bool): Boolean specifying if the sampling algo is BeamSearch
        switch_attention (bool): Boolean specifying if `Attention` is to be switched with `DecoderMaskedSelfAttention`
    """
    if is_beam_search:
        new_inputs = []
        for _i, vi in enumerate(subg.input):
            new_inputs.extend([vi])

        # Add 2 BeamSearch specific inputs
        new_inputs.extend([onnx.helper.make_tensor_value_info("beam_width", onnx.TensorProto.INT32, shape=[1])])
        new_inputs.extend(
            [
                onnx.helper.make_tensor_value_info(
                    "cache_indirection",
                    onnx.TensorProto.INT32,
                    shape=["batch_size", "beam_width", "max_seq_len"],
                )
            ]
        )
        subg.ClearField("input")
        subg.input.extend(new_inputs)

    if switch_attention:
        decoder_masked_attention_supported_attr = [
            "past_present_share_buffer",
            "num_heads",
            "scale",
            "mask_filter_value",
            "domain",
        ]

        new_nodes = []
        for node in subg.node:
            if node.op_type == "Attention":
                kwargs = kwargs_of(node)
                for k in kwargs.copy():
                    # The Attention operator does not support different qkv hidden sizes when past/present
                    # input/output exists (GPT2 model). Hence, we should never run into this.
                    # But, if we do, do not go ahead with the optimization.
                    if k == "qkv_hidden_sizes":
                        return False

                    if k not in decoder_masked_attention_supported_attr:
                        # Log the fact that we are removing certain attributes from the node
                        # We don't need to log it for "unidirectional" as we are aware that
                        # decoding attention kernels are unidirectional by definition.
                        if k != "unidirectional":
                            logger.warning(
                                f"Removing attribute: {k} from Attention node while switching to DecoderMaskedSelfAttention"
                            )

                        del kwargs[k]

                nis = []
                nis.extend(node.input)

                # Add 2 BeamSearch specific inputs
                if is_beam_search:
                    while len(nis) < 7:
                        nis.extend([""])
                    if len(nis) < 8:
                        nis.extend(["beam_width"])
                    if len(nis) < 9:
                        nis.extend(["cache_indirection"])

                node = onnx.helper.make_node(  # noqa: PLW2901
                    "DecoderMaskedSelfAttention",
                    nis,
                    node.output,
                    name=node.name,
                    **kwargs,
                )
            new_nodes.extend([node])
        subg.ClearField("node")
        subg.node.extend(new_nodes)

    return True


def find_past_seq_len_usage(subg: GraphProto):
    """Correct graph which originally use dim of past_seq_len from input_ids's shape which is fixed to max_seq_len after
       shared past/present buffer

    Args:
        subg (GraphProto): GraphProto of the decoder subgraph
    return:
        tensor_names_to_rename : set of tensor names which is equal to past_sequence_length
        nodes_to_remove : list of node to remove
    """
    tensor_names_to_rename = set()
    nodes_to_remove = []

    graph_input_names = {inp.name: index for index, inp in enumerate(subg.input)}

    input_name_to_nodes = {}
    output_name_to_node = {}
    for node in subg.node:
        for input_name in node.input:
            if input_name:
                if input_name not in input_name_to_nodes:
                    input_name_to_nodes[input_name] = [node]
                else:
                    input_name_to_nodes[input_name].append(node)
        for output_name in node.output:
            if output_name:
                output_name_to_node[output_name] = node

    for node in subg.node:
        # find "past_key_self_0 --> [Transpose(past_key_self_0) --> Reshape(past_key_self_0)] --> Shape(past_key_self_0) --> Gather(*, 2)"
        # where [Transpose(past_key_self_0) --> Reshape(past_key_self_0)] may or may not exist
        if node.op_type == "Gather":
            if not node.input[1] or not node.input[0]:
                continue

            # Find Gather node's index value
            shape_tensor_name, shape_index_name = (node.input[0], node.input[1])
            ini_gather_indices = None
            if "Constant_" in shape_index_name:
                # If shape_index_name refers to a Constant node
                for const_node in subg.node:
                    if const_node.op_type == "Constant" and const_node.output[0] == shape_index_name:
                        ini_gather_indices = const_node.attribute[0].t
                        break
            else:
                # If shape_index_name refers to an initializer
                for tensor in subg.initializer:
                    if tensor.name == shape_index_name:
                        ini_gather_indices = tensor
                        break
            if ini_gather_indices is None:
                continue
            gather_indices_arr = onnx.numpy_helper.to_array(ini_gather_indices)

            if (
                gather_indices_arr.size == 1
                and gather_indices_arr.item() in {1, 2}
                and node.input[0] in output_name_to_node
            ):
                shape_node = output_name_to_node[shape_tensor_name]
                if not (shape_node.op_type == "Shape" and shape_node.input[0]):
                    continue

                if (
                    shape_node.input[0] in graph_input_names
                    and (
                        shape_node.input[0].startswith("past_key_self_")
                        or shape_node.input[0].startswith("past_value_self_")
                    )
                    and gather_indices_arr.item() == 2
                ):
                    # "past_key_self_0 --> Shape(past_key_self_0) --> Gather(*, 2)"
                    tensor_names_to_rename.add(node.output[0])
                    nodes_to_remove.append(node)
                    if len(input_name_to_nodes[shape_node.output[0]]) == 1:
                        nodes_to_remove.append(shape_node)
                        continue

                if shape_node.input[0] not in output_name_to_node:
                    continue
                reshape_node = output_name_to_node[shape_node.input[0]]
                if not (reshape_node.op_type == "Reshape" and reshape_node.input[0]):
                    continue
                transpose_node = output_name_to_node[reshape_node.input[0]]
                if not (transpose_node.op_type == "Transpose" and transpose_node.input[0]):
                    continue

                if (
                    transpose_node.input[0] in graph_input_names
                    and (
                        transpose_node.input[0].startswith("past_key_self_")
                        or transpose_node.input[0].startswith("past_value_self_")
                    )
                    and gather_indices_arr.item() == 1
                ):
                    # "past_key_self_0 --> Transpose(past_key_self_0) --> Reshape(past_key_self_0) --> Shape(past_key_self_0) --> Gather(*, 2)"
                    tensor_names_to_rename.add(node.output[0])
                    nodes_to_remove.extend([node, shape_node, reshape_node])
                    if len(input_name_to_nodes[transpose_node.output[0]]) == 1:
                        nodes_to_remove.append(transpose_node)
                        continue

    return tensor_names_to_rename, nodes_to_remove


def add_cache_indirection_to_mha(model: OnnxModel, past_seq_len_name: str):
    # Add past_sequence_length and cache_indirection as inputs to all MultiHeadAttention ops and as inputs to model
    cache_indirection_name = "cache_indirection"
    mha_nodes = list(filter(lambda node: node.op_type == "MultiHeadAttention", model.model.graph.node))
    for node in mha_nodes:
        # MHA op takes the following potential inputs:
        # query, key, value, bias, key_padding_mask, add_qk, past_key, past_value
        while len(node.input) < 8:
            node.input.append("")
        node.input.append(past_seq_len_name)
        node.input.append(cache_indirection_name)

    model.model.graph.input.append(
        onnx.helper.make_tensor_value_info(
            cache_indirection_name, TensorProto.INT32, shape=["batch_size", "beam_width", "max_sequence_length"]
        ),
    )
    model.topological_sort()
    return model


def add_output_qk_to_mha(model: OnnxModel, dtype: int = 0, skip_node_idxs: list[int] = []):  # noqa: B006
    # Add output_qk as output to MultiHeadAttention ops and as outputs to model
    output_qk_basename = "output_cross_qk"
    output_qks = []
    mha_nodes = list(filter(lambda node: node.op_type == "MultiHeadAttention", model.model.graph.node))
    for idx, node in enumerate(mha_nodes):
        # Skip MHA nodes where output_qk does not need to be added
        if idx in skip_node_idxs:
            continue

        # Get `num_heads` attribute from MHA
        num_heads = 0
        for att in node.attribute:
            if att.name == "num_heads":
                num_heads = att.i
                break

        # Get dtype for `output_qk` based on MHA bias if not provided
        output_qk_dtype = dtype
        if output_qk_dtype == 0:
            for i in model.model.graph.initializer:
                if i.name == node.input[3]:
                    output_qk_dtype = i.data_type
                    break

        # Get `target_sequence_length` attribute from 4D input for key if it's a constant
        target_sequence_length = "target_sequence_length"
        for i in model.model.graph.input:
            if i.name == node.input[1]:
                target_sequence_length = i.type.tensor_type.shape.dim[2].dim_value
                break

        # MHA op takes the following potential outputs:
        # output, present_key, present_value
        while len(node.output) < 3:
            node.output.append("")

        output_qk_name = f"{output_qk_basename}_{idx // 2}"
        node.output.append(output_qk_name)
        output_qks.append(
            onnx.helper.make_tensor_value_info(
                output_qk_name,
                output_qk_dtype,
                shape=["batch_size", num_heads, "sequence_length", target_sequence_length],
            ),
        )

    model.model.graph.output.extend(output_qks)
    model.topological_sort()
    return model


def fix_past_sequence_length(model: OnnxModel):
    # Modify total_sequence_length = past_sequence_length + curr_sequence_length subgraph to calculate
    # past_sequence_length from the new `past_sequence_length` input of size 1D and type int32 instead of
    # from `past_key_self_0` since DecoderMaskedMultiHeadAttention (DMMHA) uses buffer sharing and
    # `past_key_self_0.shape[2] = max_sequence_length` instead of `past_key_self_0.shape[2] = past_sequence_length`
    # when buffer sharing is enabled
    #
    # Before:
    #
    #   input_ids      past_key_self_0
    #       |                 |
    #     Shape             Shape
    #       |                 |
    #     Gather            Gather
    #     (idx=1)           (idx=2)
    #       |                 |    \
    #       +--------+--------+    Unsqueeze
    #                |
    #               Add
    #
    # After:
    #
    #   input_ids    past_sequence_length (1D)
    #       |                 |
    #     Shape            Squeeze
    #       |                 |
    #     Gather             Cast
    #     (idx=1)           (int64)
    #       |                 |    \
    #       +--------+--------+    Unsqueeze
    #                |
    #               Add

    # Constant names to be used
    past_seq_len_name = "past_sequence_length"
    past_seq_len_int32 = "past_seq_len_int32"
    past_seq_len_int64 = "past_seq_len_int64"

    node = list(filter(lambda n: n.op_type == "LayerNormalization", model.model.graph.node))[0]  # noqa: RUF015

    base_path_hf = model.match_parent_path(
        node,
        ["Add", "Gather", "Tile", "Expand", "Unsqueeze", "Range"],
        [0, 1, 1, 0, 0, 0],
    )
    base_path_oai = model.match_parent_path(
        node,
        ["Add", "Slice"],
        [0, 1],
    )
    if base_path_hf is not None:
        base_path = base_path_hf
    elif base_path_oai is not None:
        base_path = base_path_oai
    else:
        logger.info("Cannot identify base path for fixing past_sequence_length subgraph")
        return
    base_node = base_path[-1]

    if base_node.op_type == "Range":
        # Hugging Face implementation
        range_node = base_path[-1]

        gather_path = model.match_parent_path(
            range_node,
            ["Gather", "Shape"],
            [0, 0],
        )
        if gather_path is None:
            logger.info("Cannot identify gather path for fixing past_sequence_length subgraph")
            return

        add_path = model.match_parent_path(
            range_node,
            ["Add", "Gather", "Shape"],
            [1, 0, 0],
        )
        if add_path is None:
            logger.info("Cannot identify add path for fixing past_sequence_length subgraph")
            return
        add_node = add_path[0]

        if gather_path != add_path[1:]:
            logger.info("Gather path and add path do not share the same nodes for calculating the past_sequence_length")
            return

        # Remove `past_key_self_0 --> Shape --> Gather` connection
        constant_in_gather = list(filter(lambda n: n.output[0] == gather_path[0].input[1], model.model.graph.node))[0]  # noqa: RUF015
        model.model.graph.node.remove(constant_in_gather)
        model.model.graph.node.remove(gather_path[0])
        model.model.graph.node.remove(gather_path[1])

        # Add `past_seq_len_int64` as an input name to existing nodes
        range_node.input[0] = past_seq_len_int64
        add_node.input[0] = past_seq_len_int64

    else:
        # OpenAI implementation
        input_ids_path = model.match_parent_path(
            base_node,
            ["Unsqueeze", "Add", "Gather", "Shape", "Reshape", "Transpose"],
            [2, 0, 0, 0, 0, 0],
        )
        if input_ids_path is None:
            logger.info("Cannot identify input_ids path for fixing past_sequence_length subgraph")
            return
        add_node = input_ids_path[1]

        past_key_path = model.match_parent_path(
            base_node,
            ["Unsqueeze", "Gather", "Shape", "Reshape", "Transpose"],
            [1, 0, 0, 0, 0],
        )
        if past_key_path is None:
            logger.info("Cannot identify past_key path for fixing past_sequence_length subgraph")
            return
        unsqueeze_node = past_key_path[0]

        if input_ids_path[2:] != past_key_path[1:]:
            logger.info(
                "The input_ids path and past_key path do not share the same nodes for calculating the past_sequence_length"
            )
            return

        # Remove `past_key_self_0 --> Transpose --> Reshape --> Shape --> Gather` connection
        constant_in_gather = list(filter(lambda n: n.output[0] == past_key_path[1].input[1], model.model.graph.node))[0]  # noqa: RUF015
        model.model.graph.node.remove(constant_in_gather)
        constant_in_reshape = list(filter(lambda n: n.output[0] == past_key_path[-2].input[1], model.model.graph.node))[  # noqa: RUF015
            0
        ]
        model.model.graph.node.remove(constant_in_reshape)
        model.model.graph.node.remove(past_key_path[1])
        model.model.graph.node.remove(past_key_path[2])
        model.model.graph.node.remove(past_key_path[3])
        model.model.graph.node.remove(past_key_path[4])

        # Add `past_seq_len_int64` as an input name to existing nodes
        unsqueeze_node.input[0] = past_seq_len_int64
        add_node.input[0] = past_seq_len_int64

    # Add `past_sequence_length` as model input
    model.model.graph.input.append(
        onnx.helper.make_tensor_value_info(past_seq_len_name, TensorProto.INT32, shape=[1]),
    )

    # Add `past_sequence_length --> Squeeze --> Cast` connection
    squeeze_node = onnx.helper.make_node(
        "Squeeze",
        inputs=[past_seq_len_name],
        outputs=[past_seq_len_int32],
        name=model.create_node_name("Squeeze"),
    )
    squeeze_output = onnx.helper.make_tensor_value_info(past_seq_len_int32, TensorProto.INT32, shape=[])
    cast_node = onnx.helper.make_node(
        "Cast",
        inputs=[past_seq_len_int32],
        outputs=[past_seq_len_int64],
        name=model.create_node_name("Cast"),
        to=TensorProto.INT64,
    )
    cast_output = onnx.helper.make_tensor_value_info(past_seq_len_int64, TensorProto.INT64, shape=[])

    # Add new nodes to graph
    model.model.graph.node.extend([squeeze_node, cast_node])
    model.model.graph.value_info.extend([squeeze_output, cast_output])
    model.topological_sort()
    return model, past_seq_len_name


def replace_mha_with_dmmha(model: OnnxModel, past_seq_len_name: str):
    # Add `beam_width` and `cache_indirection` as model inputs
    beam_width = "beam_width"
    cache_indirection = "cache_indirection"

    model.model.graph.input.extend(
        [
            onnx.helper.make_tensor_value_info(beam_width, TensorProto.INT32, shape=[1]),
            onnx.helper.make_tensor_value_info(
                cache_indirection, TensorProto.INT32, shape=["batch_size", "beam_width", "max_sequence_length"]
            ),
        ]
    )

    # Replace all `MultiHeadAttention` nodes with `DecoderMaskedMultiHeadAttention` nodes
    mha_nodes = list(filter(lambda node: node.op_type == "MultiHeadAttention", model.model.graph.node))
    for idx, node in enumerate(mha_nodes):
        # Get `num_heads` attribute from MHA
        num_heads = 0
        for att in node.attribute:
            if att.name == "num_heads":
                num_heads = att.i
                break

        # Make Q*K outputs for cross-attention layers, which happen every alternative layer
        qk_output_name = f"output_cross_qk_{idx // 2}"
        qk_output = onnx.helper.make_tensor_value_info(
            qk_output_name, TensorProto.FLOAT, shape=["batch_size", num_heads, 1, "encode_sequence_length / 2"]
        )
        if idx % 2 == 1:
            model.model.graph.output.append(qk_output)

        # Make DMMHA node
        dmmha_node = onnx.helper.make_node(
            "DecoderMaskedMultiHeadAttention",
            inputs=[
                node.input[0],  # query
                node.input[1],  # key
                node.input[2],  # value
                "",  # mask_index
                "",  # relative_position_bias
                node.input[6] if len(node.input) > 4 else "",  # past_key
                node.input[7] if len(node.input) > 4 else "",  # past_value
                past_seq_len_name,  # past_sequence_length
                beam_width,  # beam_width
                cache_indirection,  # cache_indirection
                node.input[3],  # bias
            ],
            outputs=[
                node.output[0],  # output
                node.output[1] if len(node.input) > 4 else "",  # present_key
                node.output[2] if len(node.input) > 4 else "",  # present_value
                qk_output_name if idx % 2 == 1 else "",  # output_cross_qk
            ],
            name=node.name.replace("MultiHeadAttention", "DecoderMaskedMultiHeadAttention"),
            domain="com.microsoft",
            num_heads=num_heads,
            output_qk=(idx % 2),
            past_present_share_buffer=1,
        )
        if idx % 2 == 0:
            # Remove empty string for output_cross_qk, which happens every alternative layer
            dmmha_node.output.remove("")

        model.model.graph.node.remove(node)
        model.model.graph.node.extend([dmmha_node])

    model.topological_sort()
    return model


def replace_mha_with_gqa(
    model: OnnxModel,
    attn_mask: str,
    kv_num_heads: int = 0,
    world_size: int = 1,
    window_size: int = -1,
):
    # Insert attention_mask subgraph to calculate shared inputs for all GroupQueryAttention nodes
    #
    #                attention_mask
    #               /              \
    #          ReduceSum          Shape
    #              |                |
    #             Sub             Gather
    #              |                |
    #          seqlens_k   total_sequence_length
    #              |                |
    #        Cast to int32    Cast to int32

    model.add_initializer(
        onnx.helper.make_tensor(
            name="one",
            data_type=TensorProto.INT64,
            dims=[1],
            vals=[1],
        )
    )
    reduce_sum_node = onnx.helper.make_node(
        "ReduceSum",
        inputs=[attn_mask, "one"],
        outputs=[attn_mask + "_row_sums"],
        name=model.create_node_name("ReduceSum"),
    )
    sub_node = onnx.helper.make_node(
        "Sub",
        inputs=[attn_mask + "_row_sums", "one"],
        outputs=["seqlens_k_int64"],
        name=model.create_node_name("Sub"),
    )
    seqlen_k_cast_node = onnx.helper.make_node(
        "Cast",
        inputs=["seqlens_k_int64"],
        outputs=["seqlens_k"],
        name=model.create_node_name("Cast"),
        to=TensorProto.INT32,
    )
    shape_node = onnx.helper.make_node(
        "Shape",
        inputs=[attn_mask],
        outputs=[attn_mask + "_shape"],
        name=model.create_node_name("Shape"),
    )
    gather_node = onnx.helper.make_node(
        "Gather",
        inputs=[attn_mask + "_shape", "one"],
        outputs=["total_seq_len_int64"],
        name=model.create_node_name("Gather"),
        axis=0,
    )
    total_seqlen_cast_node = onnx.helper.make_node(
        "Cast",
        inputs=["total_seq_len_int64"],
        outputs=["total_seq_len"],
        name=model.create_node_name("Cast"),
        to=TensorProto.INT32,
    )
    model.model.graph.node.extend(
        [
            reduce_sum_node,
            sub_node,
            seqlen_k_cast_node,
            shape_node,
            gather_node,
            total_seqlen_cast_node,
        ]
    )

    # Replace MultiHeadAttention with GroupQueryAttention
    #
    # When replacing, fuse the following subgraph:
    #
    #                 root_input
    #               /     |      \
    #         MatMul    MatMul    MatMul
    #           |         |         |
    #          Add       Add       Add      (optional Adds)
    #           |         |         |
    #         RotEmb    RotEmb      |
    #            \        |        /
    #             MultiHeadAttention
    #
    # to this new subgraph:
    #
    #                 root_input
    #                     |
    #                PackedMatMul           (if possible)
    #                     |
    #                 PackedAdd             (if possible)
    #                     |
    #             GroupQueryAttention
    #

    mha_nodes = list(filter(lambda node: node.op_type == "MultiHeadAttention", model.model.graph.node))
    for idx, node in enumerate(mha_nodes):
        # Detect Q path to MHA
        q_path_1 = model.match_parent_path(node, ["RotaryEmbedding", "Add", "MatMul"], [0, 0, 0])
        q_path_2 = model.match_parent_path(node, ["RotaryEmbedding", "MatMul"], [0, 0])

        q_rotary, q_add, q_matmul = None, None, None
        if q_path_1 is not None:
            q_rotary, q_add, q_matmul = q_path_1
        elif q_path_2 is not None:
            q_rotary, q_matmul = q_path_2

        # Detect K path to MHA
        k_path_1 = model.match_parent_path(node, ["RotaryEmbedding", "Add", "MatMul"], [1, 0, 0])
        k_path_2 = model.match_parent_path(node, ["RotaryEmbedding", "MatMul"], [1, 0])

        k_rotary, k_add, k_matmul = None, None, None
        if k_path_1 is not None:
            k_rotary, k_add, k_matmul = k_path_1
        elif k_path_2 is not None:
            k_rotary, k_matmul = k_path_2

        # Detect V path to MHA
        v_path_1 = model.match_parent_path(node, ["Add", "MatMul"], [2, 0])
        v_path_2 = model.match_parent_path(node, ["MatMul"], [2])

        v_add, v_matmul = None, None
        if v_path_1 is not None:
            v_add, v_matmul = v_path_1
        elif v_path_2 is not None:
            v_matmul = v_path_2[0]

        # Get `interleaved` attribute from RotaryEmbedding
        interleaved = 0
        if q_rotary is not None and k_rotary is not None:
            for att in q_rotary.attribute:
                if att.name == "interleaved":
                    interleaved = att.i

        # Get `num_heads` attribute from MHA
        num_heads = 0
        for att in node.attribute:
            if att.name == "num_heads":
                num_heads = att.i

        # Check if root_input to Q/K/V paths is the same
        root_input_is_same = q_matmul.input[0] == k_matmul.input[0] and k_matmul.input[0] == v_matmul.input[0]

        # Check if Q/K/V paths all have bias or all don't have bias
        all_paths_have_bias = q_add is not None and k_add is not None and v_add is not None
        all_paths_have_no_bias = q_add is None and k_add is None and v_add is None

        # Make PackedMatMul node if possible
        q_input_to_attention, k_input_to_attention, v_input_to_attention = "", "", ""
        if root_input_is_same and (all_paths_have_bias or all_paths_have_no_bias):
            qw = NumpyHelper.to_array(model.get_initializer(q_matmul.input[1]))
            kw = NumpyHelper.to_array(model.get_initializer(k_matmul.input[1]))
            vw = NumpyHelper.to_array(model.get_initializer(v_matmul.input[1]))

            dim = qw.shape[-1]
            qkv_weight = np.stack((qw, kw, vw), axis=1).reshape(dim, 3 * dim)
            qkv_weight = onnx.numpy_helper.from_array(qkv_weight, name=f"QKV_Weight_{idx}")
            model.add_initializer(qkv_weight)

            packed_matmul_node = onnx.helper.make_node(
                "MatMul",
                inputs=[q_matmul.input[0], qkv_weight.name],
                outputs=[f"{qkv_weight.name}_output"],
                name=model.create_node_name("MatMul"),
            )
            model.model.graph.node.extend([packed_matmul_node])
            model.model.graph.node.remove(q_matmul)
            model.model.graph.node.remove(k_matmul)
            model.model.graph.node.remove(v_matmul)
            q_input_to_attention = packed_matmul_node.output[0]

            # Make PackedAdd node if possible
            if all_paths_have_bias:
                qb = NumpyHelper.to_array(model.get_initializer(q_add.input[1]))
                kb = NumpyHelper.to_array(model.get_initializer(k_add.input[1]))
                vb = NumpyHelper.to_array(model.get_initializer(v_add.input[1]))

                dim = qb.shape[-1]
                qkv_bias = np.stack((qb, kb, vb), axis=0).reshape(3 * dim)
                qkv_bias = onnx.numpy_helper.from_array(qkv_bias, name=f"QKV_Bias_{idx}")
                model.add_initializer(qkv_bias)
                packed_add_node = onnx.helper.make_node(
                    "Add",
                    inputs=[packed_matmul_node.output[0], qkv_bias.name],
                    outputs=[f"{qkv_bias.name}_output"],
                )
                model.model.graph.node.extend([packed_add_node])
                model.model.graph.node.remove(q_add)
                model.model.graph.node.remove(k_add)
                model.model.graph.node.remove(v_add)
                q_input_to_attention = packed_add_node.output[0]

        else:
            q_input_to_attention = q_matmul.output[0]
            k_input_to_attention = k_matmul.output[0]
            v_input_to_attention = v_matmul.output[0]

        # Make GQA node
        gqa_node = onnx.helper.make_node(
            "GroupQueryAttention",
            inputs=[
                q_input_to_attention,  # query
                k_input_to_attention,  # key
                v_input_to_attention,  # value
                node.input[6],  # past_key
                node.input[7],  # past_value
                seqlen_k_cast_node.output[0],  # seqlens_k (for attention mask)
                total_seqlen_cast_node.output[0],  # total_seq_len (for attention mask)
                (q_rotary.input[2] if q_rotary is not None else ""),  # cos_cache (for rotary embeddings)
                (q_rotary.input[3] if q_rotary is not None else ""),  # sin_cache (for rotary embeddings)
            ],
            outputs=node.output,
            name=node.name.replace("MultiHeadAttention", "GroupQueryAttention"),
            domain="com.microsoft",
            num_heads=num_heads // world_size,
            kv_num_heads=(num_heads // world_size if kv_num_heads == 0 else kv_num_heads // world_size),
            local_window_size=window_size,
            do_rotary=int(q_rotary is not None and k_rotary is not None),
            rotary_interleaved=interleaved,
        )
        model.model.graph.node.remove(node)
        model.model.graph.node.extend([gqa_node])

        if q_rotary is not None:
            model.model.graph.node.remove(q_rotary)
        if k_rotary is not None:
            model.model.graph.node.remove(k_rotary)

    return model


def update_decoder_subgraph_output_cross_attention(subg: GraphProto):
    input_self_past_0 = 1
    # w/wo attention mask, w/wo hidden_state
    graph_input_names = [gi.name for gi in subg.input]
    while input_self_past_0 < 3 and not graph_input_names[input_self_past_0].startswith("past"):
        input_self_past_0 += 1
    output_self_present_0 = 1

    num_layers = (len(subg.output) - output_self_present_0) // 2
    input_cross_past_0 = 2 * num_layers + input_self_past_0
    past_key_cross_inputs = {subg.input[layer * 2 + input_cross_past_0].name: layer for layer in range(num_layers)}
    print(f"    -- past_key_cross_inputs = {past_key_cross_inputs}")

    input_past_key_cross_0_shape = shape_of(subg.input[input_cross_past_0])
    print(f"past_key_cross_0_shape is {input_past_key_cross_0_shape}")
    batch_size_dim = input_past_key_cross_0_shape[0]
    num_heads_dim = input_past_key_cross_0_shape[1]
    cross_seq_len_dim = input_past_key_cross_0_shape[2]

    num_layer_output_qk = 0
    for node in subg.node:
        if (node.op_type == "DecoderMaskedMultiHeadAttention") and (node.input[1] in past_key_cross_inputs):
            print(f"    -- add cross QK output from: node: {node.name} with output: {node.output}")
            num_layer_output_qk += 1
            layer = past_key_cross_inputs[node.input[1]]
            cross_attention_out_name = f"output_cross_qk_{layer}"
            appended_names = [""] * (3 - len(node.output))
            appended_names.append(cross_attention_out_name)
            node.output.extend(appended_names)
            node.attribute.extend([onnx.helper.make_attribute("output_qk", 1)])

            cross_attention = onnx.helper.make_tensor_value_info(
                cross_attention_out_name,
                TensorProto.FLOAT,
                [batch_size_dim, num_heads_dim, 1, cross_seq_len_dim],
            )
            subg.output.extend([cross_attention])
    if num_layer_output_qk != num_layers:
        raise ValueError(f"Did not add cross QK for all layers{num_layers} vs {num_layer_output_qk}")


def update_decoder_subgraph_share_buffer_and_use_decoder_masked_mha(subg: ModelProto):
    input_self_past_0 = 1
    # w/wo attention mask, w/wo hidden_state
    graph_input_names = [gi.name for gi in subg.input]
    while input_self_past_0 < 3 and not graph_input_names[input_self_past_0].startswith("past"):
        input_self_past_0 += 1
    output_self_past_0 = 1

    num_layers = int((len(subg.input) - input_self_past_0) / 4)
    input_cross_past_0 = 2 * num_layers + input_self_past_0

    new_nodes = []
    old_nodes = []
    for node in subg.node:
        if node.op_type == "MultiHeadAttention":
            old_nodes.extend([node])

    # If not all the MultiHeadAttention nodes are fused, this optimization is not applicable
    if len(old_nodes) < num_layers:
        return False

    # Redirect the RelativePositionBias node's input from past_key_self_0.shape[2] to past_sequence_length.
    # There is only one RelativePositionBias node in T5 decoder subgraph.
    rel_pos_bias_node = None
    for node in subg.node:
        if node.op_type == "RelativePositionBias":
            rel_pos_bias_node = node
            break

    decoder_masked_attention_supported_attr = [
        "past_present_share_buffer",
        "num_heads",
        "scale",
        "mask_filter_value",
        "domain",
    ]

    target_squeezed_past_seq_name = "past_sequence_length_squeezed_int64"
    tensor_names_to_rename, nodes_to_remove = find_past_seq_len_usage(subg)
    if len(tensor_names_to_rename) > 0:
        for name_to_rename in tensor_names_to_rename:
            print(f"Found tensor name `{name_to_rename}` to be renamed to `{target_squeezed_past_seq_name}`")
        for nr in nodes_to_remove:
            print(f"Found node to remove: type = {nr.op_type}, name = {nr.name}")

        squeeze_node = onnx.helper.make_node(
            "Squeeze",
            ["past_sequence_length"],
            ["past_sequence_length_squeezed"],
            name="node_past_sequence_length_squeeze",
        )
        cast_node = onnx.helper.make_node(
            "Cast",
            ["past_sequence_length_squeezed"],
            [target_squeezed_past_seq_name],
            name="node_past_sequence_length_squeeze_cast",
            to=TensorProto.INT64,
        )
        new_nodes.extend([squeeze_node, cast_node])

    for node in subg.node:
        if len(node.output) > 0 and rel_pos_bias_node is not None and node.output[0] == rel_pos_bias_node.input[1]:
            cast_node = onnx.helper.make_node(
                "Cast",
                ["past_sequence_length"],
                ["past_sequence_length_int64"],
                name="past_sequence_length_cast",
                to=TensorProto.INT64,
            )
            node.input[1] = cast_node.output[0]
            new_nodes.extend([cast_node])

        if node.op_type == "MultiHeadAttention":
            kwargs = kwargs_of(node)
            for k in kwargs.copy():
                if k not in decoder_masked_attention_supported_attr:
                    del kwargs[k]

            # note: This logic only apply to T5 model where there is no bias in Attention node.
            nis = [
                node.input[0],  # query
                node.input[1],  # key
                node.input[2],  # value
            ]

            nis.extend([node.input[4] if len(node.input) > 4 else ""])  # 2D mask
            nis.extend([node.input[5] if len(node.input) > 5 else ""])  # attention_bias
            nis.extend([node.input[6] if len(node.input) > 6 else ""])  # past_key
            nis.extend([node.input[7] if len(node.input) > 7 else ""])  # past_value
            nis.extend(["past_sequence_length"])  # past_sequence_length
            nis.extend(["beam_width"])  # beam_width
            nis.extend(["cache_indirection"])  # cache_indirection
            nis.extend([node.input[3] if len(node.input) > 3 else ""])  # bias

            kwargs["past_present_share_buffer"] = 1

            node = onnx.helper.make_node(  # noqa: PLW2901
                "DecoderMaskedMultiHeadAttention",
                nis,
                node.output,
                name=node.name,
                **kwargs,
            )

        if node not in nodes_to_remove:
            for index, name in enumerate(node.input):
                if name in tensor_names_to_rename:
                    node.input[index] = target_squeezed_past_seq_name
            new_nodes.extend([node])

    subg.ClearField("node")
    subg.node.extend(new_nodes)
    orig_input_names = [inp.name for inp in subg.input]

    new_inputs = []
    for i, vi in enumerate(subg.input):
        if i >= input_self_past_0 and i < input_cross_past_0:
            shape = shape_of(vi)
            vi = onnx.helper.make_tensor_value_info(  # noqa: PLW2901
                vi.name,
                elem_type=vi.type.tensor_type.elem_type,
                shape=[shape[0], shape[1], "max_seq_len", shape[3]],
            )
        new_inputs.extend([vi])
    if "past_sequence_length" not in orig_input_names:
        new_inputs.extend(
            [onnx.helper.make_tensor_value_info("past_sequence_length", onnx.TensorProto.INT32, shape=[1])]
        )
    if "beam_width" not in orig_input_names:
        new_inputs.extend([onnx.helper.make_tensor_value_info("beam_width", onnx.TensorProto.INT32, shape=[1])])
    if "cache_indirection" not in orig_input_names:
        new_inputs.extend(
            [
                onnx.helper.make_tensor_value_info(
                    "cache_indirection",
                    onnx.TensorProto.INT32,
                    shape=["batch_size", "beam_width", "max_seq_len"],
                )
            ]
        )
    subg.ClearField("input")
    subg.input.extend(new_inputs)

    new_outputs = []
    for i, vi in enumerate(subg.output):
        if i >= output_self_past_0:
            shape = shape_of(vi)
            vi = onnx.helper.make_tensor_value_info(  # noqa: PLW2901
                vi.name,
                elem_type=vi.type.tensor_type.elem_type,
                shape=[shape[0], shape[1], "max_seq_len", shape[3]],
            )
        new_outputs.extend([vi])
    subg.ClearField("output")
    subg.output.extend(new_outputs)

    return True


def pack_qkv_for_decoder_masked_mha(model_proto: ModelProto):
    onnx_model = OnnxModel(model_proto)
    output_name_to_node = onnx_model.output_name_to_node()

    nodes_to_add = []
    nodes_to_remove = []
    for node in onnx_model.nodes():
        if node.op_type == "DecoderMaskedMultiHeadAttention":
            if "past_key_cross" in node.input[1] and "past_value_cross" in node.input[2]:
                continue
            q_matmul = output_name_to_node[node.input[0]]
            k_matmul = output_name_to_node[node.input[1]]
            v_matmul = output_name_to_node[node.input[2]]

            q_weight = onnx_model.get_initializer(q_matmul.input[1])
            k_weight = onnx_model.get_initializer(k_matmul.input[1])
            v_weight = onnx_model.get_initializer(v_matmul.input[1])
            if not (q_weight and k_weight and v_weight):
                return False

            qw = NumpyHelper.to_array(q_weight)
            kw = NumpyHelper.to_array(k_weight)
            vw = NumpyHelper.to_array(v_weight)

            qkv_weight = np.concatenate([qw, kw, vw], axis=1)

            matmul_node_name = onnx_model.create_node_name("MatMul", name_prefix="MatMul_QKV")
            weight = onnx.helper.make_tensor(
                name=matmul_node_name + "_weight",
                data_type=(TensorProto.FLOAT if q_weight.data_type == 1 else TensorProto.FLOAT16),
                dims=[qkv_weight.shape[0], qkv_weight.shape[1]],
                vals=qkv_weight.flatten().tolist(),
            )

            model_proto.graph.initializer.extend([weight])

            matmul_node = onnx.helper.make_node(
                "MatMul",
                inputs=[q_matmul.input[0], matmul_node_name + "_weight"],
                outputs=[matmul_node_name + "_out"],
                name=matmul_node_name,
            )

            node.input[0] = matmul_node.output[0]
            node.input[1] = ""
            node.input[2] = ""

            nodes_to_add.extend([matmul_node])
            nodes_to_remove.extend([q_matmul, k_matmul, v_matmul])

    onnx_model.add_nodes(nodes_to_add)
    onnx_model.remove_nodes(nodes_to_remove)
    onnx_model.update_graph()

    onnx_model.topological_sort()

    return True


def update_input_shapes_for_gpt2_decoder_model(decoder_onnx_path: str, use_external_data_format: bool = True):
    """Update the input shapes for the inputs "input_ids" and "position_ids" and make the sequence length dim value 1 for each of them.
       The decoder model will be over-written.

    Args:
        decoder_onnx_path (str): Path of GPT-2 decoder onnx model
        use_external_data_format(bool): output tensors to external data or not.
    """

    decoder_model_proto = onnx.load_model(decoder_onnx_path, load_external_data=True)
    for i in range(len(decoder_model_proto.graph.input)):
        if (
            decoder_model_proto.graph.input[i].name == "input_ids"
            or decoder_model_proto.graph.input[i].name == "position_ids"
        ):
            shape_dim_proto = decoder_model_proto.graph.input[i].type.tensor_type.shape.dim[1]

            # Clear any existing dim_param first
            if shape_dim_proto.HasField("dim_param"):
                shape_dim_proto.Clear()

            # Update dim_value to be 1
            shape_dim_proto.dim_value = 1

    OnnxModel.save(
        decoder_model_proto,
        decoder_onnx_path,
        save_as_external_data=use_external_data_format,
    )
    return True


def generate_gpt2_init_decoder(
    decoder_onnx_path: str,
    init_decoder_onnx_path: str,
    use_external_data_format: bool = True,
) -> bool:
    """Generates the initial decoder GPT2 subgraph and saves it for downstream use.
       The initial decoder model will be saved to init_decoder_onnx_path.

    Args:
        decoder_onnx_path (str): Path of GPT-2 decoder onnx model
        init_decoder_onnx_path (str): Path of GPT-2 init decoder onnx model
        use_external_data_format(bool): output tensors to external data or not.
    """
    init_decoder_model_proto = onnx.load_model(decoder_onnx_path, load_external_data=True)

    logits_output_name = init_decoder_model_proto.graph.output[0].name

    gpt2_init_decoder_model = OnnxModel(init_decoder_model_proto)

    output_name_to_node = gpt2_init_decoder_model.output_name_to_node()
    assert logits_output_name in output_name_to_node

    logits_matmul_node = output_name_to_node[logits_output_name]

    # Sanity check - the logits need to be produced by a MatMul node
    if logits_matmul_node.op_type != "MatMul":
        return False

    # Try to find the last residual Add
    # For fp16, there are Casts along the way

    # Normalization Node is : LayerNormalization
    logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
        logits_matmul_node,
        [
            "Cast",
            "LayerNormalization",
            "Add",
            "Add",
            "Cast",
            "MatMul",
            "Cast",
            "FastGelu",
            "Cast",
            "MatMul",
            "Cast",
            "LayerNormalization",
            "Add",
        ],
        [0, 0, 0, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0],
    )

    # Normalization Node is : SkipLayerNormalization
    if logits_matmul_to_residual_add_path is None:
        logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
            logits_matmul_node,
            [
                "Cast",
                "SkipLayerNormalization",
                "Cast",
                "MatMul",
                "Cast",
                "FastGelu",
                "Cast",
                "MatMul",
                "Cast",
                "SkipLayerNormalization",
            ],
            [0, 0, 1, 0, 0, 0, 0, 0, 0, 0],
        )

    # Try without the Casts before and after the MatMuls
    if logits_matmul_to_residual_add_path is None:
        # Normalization Node is : LayerNormalization
        logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
            logits_matmul_node,
            [
                "LayerNormalization",
                "Add",
                "Add",
                "MatMul",
                "FastGelu",
                "MatMul",
                "LayerNormalization",
                "Add",
            ],
            [0, 0, 1, 0, 0, 0, 0, 0],
        )

        # Normalization Node is : SkipLayerNormalization
        if logits_matmul_to_residual_add_path is None:
            logits_matmul_to_residual_add_path = gpt2_init_decoder_model.match_parent_path(
                logits_matmul_node,
                [
                    "SkipLayerNormalization",
                    "MatMul",
                    "FastGelu",
                    "MatMul",
                    "SkipLayerNormalization",
                ],
                [0, 1, 0, 0, 0],
            )

    # TODO(hasesh): Are there more permutations to try before returning ?
    if logits_matmul_to_residual_add_path is None:
        return False

    residual_add_node = logits_matmul_to_residual_add_path[-1]

    # If the last node in the pattern is SkipLayerNormalization, we need to adjust our pattern searches accordingly
    is_skiplayernorm_path = residual_add_node.op_type == "SkipLayerNormalization"

    # Regular LayerNormalization path
    if not is_skiplayernorm_path:
        residual_add_to_attention_parent_index = 0
        residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
            residual_add_node,
            ["Add", "Cast", "MatMul", "Attention"],
            [residual_add_to_attention_parent_index, 0, 0, 0],
        )

        # Try other parent index of the residual Add node
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 1
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["Add", "Cast", "MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0, 0, 0],
            )

        # Try without the Casts before and after the MatMuls
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 0
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["Add", "MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0, 0],
            )

        # Try without the Casts before and after the MatMuls and other parent index of the residual Add node
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 1
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["Add", "MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0, 0],
            )

    # SkipLayerNormalization path
    else:
        residual_add_to_attention_parent_index = 0
        residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
            residual_add_node,
            ["Cast", "MatMul", "Attention"],
            [residual_add_to_attention_parent_index, 0, 0],
        )

        # Try other parent index of the residual Add node
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 1
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["Cast", "MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0, 0],
            )

        # Try without the Casts before and after the MatMuls
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 0
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0],
            )

        # Try without the Casts before and after the MatMuls and other parent index of the residual Add node
        if residual_add_to_attention_path is None:
            residual_add_to_attention_parent_index = 1
            residual_add_to_attention_path = gpt2_init_decoder_model.match_parent_path(
                residual_add_node,
                ["MatMul", "Attention"],
                [residual_add_to_attention_parent_index, 0],
            )

    # TODO(hasesh): Are there more permutations to try before returning ?
    if residual_add_to_attention_path is None:
        return False

    residual_add_to_add_parent_index = 0 if residual_add_to_attention_parent_index == 1 else 1

    # Regular LayerNormalization path
    if not is_skiplayernorm_path:
        add_before_residual_add = gpt2_init_decoder_model.match_parent(
            residual_add_node, "Add", residual_add_to_add_parent_index
        )

    # SkipLayerNormalization path
    else:
        add_before_residual_add = gpt2_init_decoder_model.match_parent(
            residual_add_node,
            "SkipLayerNormalization",
            residual_add_to_add_parent_index,
        )

    if add_before_residual_add is None:
        return False

    attention = residual_add_to_attention_path[-1]
    matmul_after_attention = residual_add_to_attention_path[-2]

    slice_starts = onnx.helper.make_tensor(
        name="SliceLastTokenStarts",
        data_type=TensorProto.INT32,
        dims=[1],
        vals=[-1],
    )

    slice_ends = onnx.helper.make_tensor(
        name="SliceLastTokenEnds",
        data_type=TensorProto.INT32,
        dims=[1],
        vals=[-2],
    )

    slice_axes = onnx.helper.make_tensor(
        name="SliceLastTokenAxes",
        data_type=TensorProto.INT32,
        dims=[1],
        vals=[1],
    )

    slice_steps = onnx.helper.make_tensor(
        name="SliceLastTokenSteps",
        data_type=TensorProto.INT32,
        dims=[1],
        vals=[-1],
    )

    gpt2_init_decoder_model.add_initializer(slice_starts)
    gpt2_init_decoder_model.add_initializer(slice_ends)
    gpt2_init_decoder_model.add_initializer(slice_axes)
    gpt2_init_decoder_model.add_initializer(slice_steps)

    # Add Slice node to the graph such that it consumes the output of Attention
    slice_0_output_name = "edge_modified_" + attention.output[0]
    slice_node_0 = onnx.helper.make_node(
        "Slice",
        inputs=[
            attention.output[0],
            "SliceLastTokenStarts",
            "SliceLastTokenEnds",
            "SliceLastTokenAxes",
            "SliceLastTokenSteps",
        ],
        outputs=[slice_0_output_name],
        name=gpt2_init_decoder_model.create_node_name("Slice", "GatherLastToken_0_"),
    )

    # Add Slice node to the graph such that it consumes the output of Add before the residual Add
    # If the 'Add' output is produced by a SkipLayerNormalization node, then adjust its output
    # index appropriately
    add_before_residual_add_output = (
        add_before_residual_add.output[0] if not is_skiplayernorm_path else add_before_residual_add.output[3]
    )

    slice_1_output_name = "edge_modified_" + add_before_residual_add.output[0]
    slice_node_1 = onnx.helper.make_node(
        "Slice",
        inputs=[
            add_before_residual_add_output,
            "SliceLastTokenStarts",
            "SliceLastTokenEnds",
            "SliceLastTokenAxes",
            "SliceLastTokenSteps",
        ],
        outputs=[slice_1_output_name],
        name=gpt2_init_decoder_model.create_node_name("Slice", "GatherLastToken_1_"),
    )

    # Add the 2 Slice nodes
    gpt2_init_decoder_model.add_node(slice_node_0)
    gpt2_init_decoder_model.add_node(slice_node_1)

    # Adjust the input(s) to the nodes consuming the outputs of the added Slice nodes
    gpt2_init_decoder_model.replace_node_input(matmul_after_attention, attention.output[0], slice_0_output_name)
    gpt2_init_decoder_model.replace_node_input(residual_add_node, add_before_residual_add_output, slice_1_output_name)

    # Topologically sort the updated graph
    gpt2_init_decoder_model.topological_sort()

    # Save the init decoder model
    OnnxModel.save(
        init_decoder_model_proto,
        init_decoder_onnx_path,
        save_as_external_data=use_external_data_format,
    )
    return True


def make_dim_proto_numeric_t5(model, config):
    """Make dim_proto numeric.

    Args:
        model: T5 encoder and decoder model.
        config: T5 config.
    """
    sequence_length = str(1)
    num_heads = str(config.num_heads)
    hidden_size = str(config.d_model)
    head_size = str(config.d_kv)

    for tensor in model.graph.output:
        for dim_proto in tensor.type.tensor_type.shape.dim:
            if dim_proto.HasField("dim_param") and dim_proto.dim_param in [
                sequence_length,
                num_heads,
                hidden_size,
                head_size,
            ]:
                dim_value = int(dim_proto.dim_param)
                dim_proto.Clear()
                dim_proto.dim_value = dim_value

    for tensor in model.graph.input:
        for dim_proto in tensor.type.tensor_type.shape.dim:
            if dim_proto.HasField("dim_param") and dim_proto.dim_param in [
                sequence_length,
                num_heads,
                hidden_size,
                head_size,
            ]:
                dim_value = int(dim_proto.dim_param)
                dim_proto.Clear()
                dim_proto.dim_value = dim_value


def convert_generation_model(
    args: argparse.Namespace,
    generation_type: GenerationType = GenerationType.BEAMSEARCH,
):
    """Convert model according to command line arguments.

    Args:
        args (argparse.Namespace): arguments parsed from command line
    """
    is_gpt2: bool = args.model_type == "gpt2"
    is_beamsearch: bool = generation_type == GenerationType.BEAMSEARCH
    is_greedysearch: bool = generation_type == GenerationType.GREEDYSEARCH
    is_sampling: bool = generation_type == GenerationType.SAMPLING
    past_present_share_buffer: bool = args.past_present_share_buffer

    logger.info(f"**** past_present_share_buffer={past_present_share_buffer}")
    if len(args.op_block_list) == 1 and args.op_block_list[0] == "auto":
        if is_gpt2 and args.precision == Precision.FLOAT16.value:
            args.op_block_list = [
                "Add",
                "LayerNormalization",
                "SkipLayerNormalization",
                "FastGelu",
            ]
            logger.info(f"**** Setting op_block_list to {args.op_block_list}")
            logger.info("**** use --op_block_list if you want to override the block operator list.")
        else:
            args.op_block_list = []

    if is_greedysearch or is_sampling:
        if not is_gpt2:
            raise NotImplementedError("Currently only gpt2 with greedy search/sampling is supported")
        if args.output_sequences_scores:
            raise NotImplementedError("output_sequences_scores currently is not supported in greedy search/sampling")
        if args.output_token_scores:
            raise NotImplementedError("output_token_scores currently is not supported in greedy search/sampling")

    # For BeamSearch, sharing buffers for past and present states is only supported
    # when using `use_decoder_masked_attention`
    if past_present_share_buffer and is_beamsearch and not args.use_decoder_masked_attention:
        raise ValueError(
            "`use_decoder_masked_attention` MUST be turned on to use `past_present_share_buffer` in case of BeamSearch"
        )

    # For any kind of sampling, using decoder masked multihead attention is only supported
    # when using `past_present_share_buffer`
    if args.use_decoder_masked_attention and not past_present_share_buffer:
        raise ValueError("`past_present_share_buffer` MUST be turned on to use `use_decoder_masked_attention`")

    # For any kind of sampling, using decoder masked multihead attention is only supported
    # on GPUs
    if args.use_decoder_masked_attention and not args.use_gpu:
        raise ValueError("`use_decoder_masked_attention` option is only supported on GPUs")

    if is_gpt2:
        if args.decoder_onnx and os.path.exists(args.decoder_onnx):
            logger.info(f"skip convert_to_onnx since path existed: {args.decoder_onnx}")
        else:
            if not args.decoder_onnx:
                onnx_filename = f"{args.model_name_or_path}_past_{args.precision}.onnx"
                args.decoder_onnx = Path(Path(args.output).parent, onnx_filename).as_posix()

            logger.info(f"Convert GPT model {args.model_name_or_path} to onnx {args.decoder_onnx} ...")
            gpt2_to_onnx(args)
    else:  # t5 or mt5
        if args.decoder_onnx and args.encoder_decoder_init_onnx:
            logger.info(
                f"skip convert_to_onnx since paths specified: {args.decoder_onnx} and {args.encoder_decoder_init_onnx}"
            )
        else:
            logger.info(f"Convert model {args.model_name_or_path} to onnx ...")
            t5_to_onnx(args)

    # We only want to pad the logits MatMul weight in the decoder for fp16 models.
    # The inherent assumption is that fp16 models run on GPU for which all
    # dims need to be a multiple of 8 to leverage tensor cores.
    # NOTE: We currently only support padding the MatMul logits weight for GPT2 GreedySearch/BeamSearch.
    # This can be expanded to other models/decoding strategies later
    logits_matmul_weight_padded = False
    if (
        not args.disable_pad_vocab_size
        and args.precision == Precision.FLOAT16.value
        and is_gpt2
        and (is_beamsearch or is_greedysearch or is_sampling)
    ):
        logger.info(
            f"Pad logits MatMul weights for optimal MatMul perf in fp16 on {args.decoder_onnx}. "
            "The file will be overwritten."
        )
        logits_matmul_weight_padded = pad_weights_of_logits_matmul(args.decoder_onnx, args.use_external_data_format)
        if not logits_matmul_weight_padded:
            logger.warning(
                "Tried and failed to pad logits MatMul weights. Performance may be sub-optimal for this MatMul"
            )

    gpt2_init_decoder_generated = False
    gpt2_init_decoder_onnx_path = None
    if (
        not args.disable_separate_gpt2_decoder_for_init_run
        and is_gpt2
        and (is_beamsearch or is_greedysearch or is_sampling)
    ):
        logger.info(f"Creating an initial run GPT2 decoder from {args.decoder_onnx}. ")

        gpt2_init_decoder_onnx_filename = f"gpt2_init_past_{args.precision}.onnx"

        gpt2_init_decoder_onnx_path = Path(Path(args.output).parent, gpt2_init_decoder_onnx_filename).as_posix()

        gpt2_init_decoder_generated = generate_gpt2_init_decoder(
            args.decoder_onnx,
            gpt2_init_decoder_onnx_path,
            args.use_external_data_format,
        )

        if not gpt2_init_decoder_generated:
            logger.warning(
                "Tried and failed to generate the init decoder GPT2 model. "
                "Performance may be sub-optimal for the initial decoding run"
            )

        # Update the graph input shapes for the non-initial decoder model to account
        # for the fact that the sequence length will always be 1
        if gpt2_init_decoder_generated and not update_input_shapes_for_gpt2_decoder_model(
            args.decoder_onnx, args.use_external_data_format
        ):
            # Can't proceed further - better to raise an exception
            raise ValueError("Could not update the input shapes for the non-initial decoder subgraph.")

    # If the user explicitly requests running shape inference or if we padded/mutated
    # weight(s)/input shape(s) in the decoder, we want to run shape inference to capture the new
    # shapes
    if logits_matmul_weight_padded or args.run_shape_inference or gpt2_init_decoder_generated:
        logger.info(f"Run symbolic shape inference on {args.decoder_onnx}. The file will be overwritten.")
        shape_inference(args.decoder_onnx, args.use_external_data_format)
        if gpt2_init_decoder_generated:
            logger.info(f"Run symbolic shape inference on {gpt2_init_decoder_onnx_path}. The file will be overwritten.")
            shape_inference(gpt2_init_decoder_onnx_path, args.use_external_data_format)

    if is_gpt2:
        config = GPT2Config.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
    elif args.model_type == "t5":
        config = T5Config.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
    else:
        config = MT5Config.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)

    if args.verbose:
        logger.info(f"Config={config}")

    eos_token_id = config.eos_token_id
    pad_token_id = config.eos_token_id if is_gpt2 else config.pad_token_id
    vocab_size = config.vocab_size

    # if vocab_size is given in parameters use that.
    if args.vocab_size != -1:
        vocab_size = args.vocab_size

    if args.eos_token_id != -1:
        eos_token_id = args.eos_token_id
    if args.pad_token_id != -1:
        pad_token_id = args.pad_token_id

    decoder_model = onnx.load_model(args.decoder_onnx, load_external_data=True)
    decoder_model.graph.name = f"{args.model_type} decoder"

    gpt2_init_decoder_model = None
    if args.model_type == "gpt2":
        verify_gpt2_subgraph(decoder_model.graph, args.precision)

        # If we generated the init decoder model, verify that as well
        if gpt2_init_decoder_generated:
            gpt2_init_decoder_model = onnx.load_model(gpt2_init_decoder_onnx_path, load_external_data=True)
            gpt2_init_decoder_model.graph.name = f"{args.model_type} init decoder"
            verify_gpt2_subgraph(gpt2_init_decoder_model.graph, args.precision)
    else:
        verify_t5_decoder_subgraph(decoder_model.graph, args.precision)

    inputs = None
    if is_beamsearch:
        inputs = [
            "input_ids",
            "max_length",
            "min_length",
            "num_beams",
            "num_return_sequences",
            "length_penalty",
            "repetition_penalty",
        ]
    elif is_greedysearch or is_sampling:
        inputs = [
            "input_ids",
            "max_length",
            "min_length",
            "repetition_penalty",
        ]

    if args.vocab_mask:
        inputs.append("vocab_mask")
    else:
        inputs.append("")

    if args.prefix_vocab_mask:
        inputs.append("prefix_vocab_mask")
    else:
        inputs.append("")

    if args.custom_attention_mask:
        inputs.append("attention_mask")
    else:
        inputs.append("")

    if is_sampling:
        if args.custom and args.presence_mask:
            inputs.append("presence_mask")
        else:
            inputs.append("")

        if args.seed:
            inputs.append("seed")

    outputs = ["sequences"]
    if args.output_sequences_scores:
        outputs.append("sequences_scores")

    if args.output_token_scores:
        assert args.output_sequences_scores, "--output_token_scores requires --output_sequences_scores"
        outputs.append("scores")

    node = None
    if is_beamsearch:
        node = onnx.helper.make_node(
            "BeamSearch",
            inputs=inputs,
            outputs=outputs,
            name=f"BeamSearch_{args.model_type}",
        )
    elif is_greedysearch:
        node = onnx.helper.make_node(
            "GreedySearch",
            inputs=inputs,
            outputs=outputs,
            name=f"GreedySearch_{args.model_type}",
        )
    elif is_sampling:
        node = onnx.helper.make_node(
            "Sampling",
            inputs=inputs,
            outputs=outputs,
            name=f"Sampling_{args.model_type}",
        )

    node.domain = "com.microsoft"

    attr_to_extend = None
    if is_beamsearch:
        attr_to_extend = [
            onnx.helper.make_attribute("eos_token_id", eos_token_id),
            onnx.helper.make_attribute("pad_token_id", pad_token_id),
            onnx.helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
            onnx.helper.make_attribute("early_stopping", 1 if args.early_stopping else 0),
            onnx.helper.make_attribute("model_type", 0 if args.model_type == "gpt2" else 1),
        ]
    elif is_greedysearch:
        attr_to_extend = [
            onnx.helper.make_attribute("eos_token_id", eos_token_id),
            onnx.helper.make_attribute("pad_token_id", pad_token_id),
            onnx.helper.make_attribute("model_type", 0 if args.model_type == "gpt2" else 1),
            onnx.helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
        ]
    elif is_sampling:
        attr_to_extend = [
            onnx.helper.make_attribute("eos_token_id", eos_token_id),
            onnx.helper.make_attribute("pad_token_id", pad_token_id),
            onnx.helper.make_attribute("model_type", 0 if args.model_type == "gpt2" else 1),
            onnx.helper.make_attribute("no_repeat_ngram_size", args.no_repeat_ngram_size),
            onnx.helper.make_attribute("temperature", args.temperature),
            onnx.helper.make_attribute("top_p", args.top_p),
            onnx.helper.make_attribute("filter_value", args.filter_value),
            onnx.helper.make_attribute("min_tokens_to_keep", args.min_tokens_to_keep),
            onnx.helper.make_attribute("custom", args.custom),
            onnx.helper.make_attribute("presence_penalty", args.presence_penalty),
        ]

    # Explicitly pass in the vocab size via an attribute
    if logits_matmul_weight_padded:
        attr_to_extend.extend([onnx.helper.make_attribute("vocab_size", vocab_size)])

    node.attribute.extend(attr_to_extend)

    initializers = []

    if args.model_type in ["t5", "mt5"]:
        if args.run_shape_inference:
            logger.info(f"Symbolic shape inference on {args.encoder_decoder_init_onnx}. The file will be overwritten.")
            shape_inference(args.encoder_decoder_init_onnx, args.use_external_data_format)
        encoder_model = onnx.load_model(args.encoder_decoder_init_onnx, load_external_data=True)
        suffix = "encoder" if len(encoder_model.graph.input) == 2 else "encoder and decoder init"
        encoder_model.graph.name = f"{args.model_type} {suffix}"
        verify_t5_encoder_decoder_init_subgraph(encoder_model.graph, args.precision)

        make_dim_proto_numeric_t5(encoder_model, config)
        make_dim_proto_numeric_t5(decoder_model, config)

        # Update decoder subgraph in preparation to use past present share buffer
        if past_present_share_buffer:
            if not args.use_decoder_masked_attention:
                raise ValueError("past_present_share_buffer is only supported with use_decoder_masked_attention")

            logger.info(
                "*****update t5 decoder subgraph to share past/present buffer and use decoder_masked_multihead_attention*****"
            )
            if update_decoder_subgraph_share_buffer_and_use_decoder_masked_mha(decoder_model.graph):
                logger.info("*****update t5 decoder subgraph successfully!!!*****")
            else:
                logger.info("*****DecoderMaskedMultiHeadAttention is not applied to T5 decoder*****")

            if pack_qkv_for_decoder_masked_mha(decoder_model):
                logger.info("*****pack qkv for decoder masked mha successfully!!!*****")
            else:
                logger.info("*****pack qkv for decoder masked mha failed!!!*****")

        if not args.disable_shared_initializers:
            # Unique shared initializers from the decoder and decoder_init could reduce memory usage in inference.
            initializers = get_shared_initializers(encoder_model, decoder_model)
            logger.info(
                f"{len(initializers)} shared initializers ({[i.name for i in initializers]}) in encoder and decoder subgraphs are moved to the main graph"
            )

            # TODO(tianleiwu): investigate the following which causes error in inference
            # Move initializer from subgraph to main graph could reduce memory usage in inference.
            # moved_initializers = move_initializers(encoder_model.graph)
            # logger.info(
            #     f"{len(moved_initializers)} initializers ({[i.name for i in moved_initializers]}) from the encoder are moved to the main graph"
            # )
            # initializers.extend(moved_initializers)

        assert config.decoder_start_token_id >= 0, "decoder_start_token_id should be >= 0"

        node.attribute.extend(
            [
                onnx.helper.make_attribute("encoder", encoder_model.graph),
                onnx.helper.make_attribute("decoder", decoder_model.graph),
                onnx.helper.make_attribute("decoder_start_token_id", config.decoder_start_token_id),
            ]
        )
    else:
        if gpt2_init_decoder_generated:
            # Move shared initializers (shared between init decoder and decoder models) to the main
            # graph and remove them from these models
            if not args.disable_shared_initializers:
                # Unique shared initializers from the decoder and decoder_init could reduce memory usage in inference.
                initializers = get_shared_initializers(gpt2_init_decoder_model, decoder_model)
                logger.info(
                    f"{len(initializers)} shared initializers ({[i.name for i in initializers]}) in decoder and init decoder subgraphs are moved to the main graph"
                )

            # Update init decoder subgraph in preparation to use past present share buffer
            if past_present_share_buffer:
                logger.info("*****update init decoder subgraph to make past and present share buffer******************")
                update_decoder_subgraph_past_present_share_buffer(gpt2_init_decoder_model.graph)

            # Update init decoder subgraph in preparation to use DecoderMaskedSelfAttention
            # NOTE: Even if we will not use DecoderMaskedSelfAttention in the init decoder subgraph
            # it makes the runtime changes cleaner if we keep both the init decoder and decoder subgraphs
            # same in terms of the subgraph inputs.
            if args.use_decoder_masked_attention and not update_decoder_subgraph_use_decoder_masked_attention(
                gpt2_init_decoder_model.graph, is_beamsearch, False
            ):
                raise ValueError("Could not update the init decoder subgraph to use DecoderMaskedSelfAttention")

            node.attribute.append(onnx.helper.make_attribute("init_decoder", gpt2_init_decoder_model.graph))
        else:
            # Move initializer from subgraph to main graph could reduce memory usage in inference.
            initializers = move_initializers(decoder_model.graph)
            logger.info(f"{len(initializers)} initializers from the decoder are moved to the main graph")

        # Update decoder subgraph in preparation to use past present share buffer
        if past_present_share_buffer:
            logger.info("*****update decoder subgraph to make past and present share buffer******************")
            update_decoder_subgraph_past_present_share_buffer(decoder_model.graph)

        # Update decoder subgraph in preparation to use DecoderMaskedSelfAttention
        if args.use_decoder_masked_attention and not update_decoder_subgraph_use_decoder_masked_attention(
            decoder_model.graph, is_beamsearch, True
        ):
            raise ValueError("Could not update the decoder subgraph to use DecoderMaskedSelfAttention")

        node.attribute.append(onnx.helper.make_attribute("decoder", decoder_model.graph))

    # graph inputs
    input_ids = onnx.helper.make_tensor_value_info("input_ids", TensorProto.INT32, ["batch_size", "sequence_length"])
    max_length = onnx.helper.make_tensor_value_info("max_length", TensorProto.INT32, [1])
    min_length = onnx.helper.make_tensor_value_info("min_length", TensorProto.INT32, [1])
    num_beams = onnx.helper.make_tensor_value_info("num_beams", TensorProto.INT32, [1])
    num_return_sequences = onnx.helper.make_tensor_value_info("num_return_sequences", TensorProto.INT32, [1])
    length_penalty = onnx.helper.make_tensor_value_info("length_penalty", TensorProto.FLOAT, [1])
    repetition_penalty = onnx.helper.make_tensor_value_info("repetition_penalty", TensorProto.FLOAT, [1])

    graph_inputs = None
    if is_beamsearch:
        graph_inputs = [
            input_ids,
            max_length,
            min_length,
            num_beams,
            num_return_sequences,
            length_penalty,
            repetition_penalty,
        ]
    elif is_greedysearch or is_sampling:
        graph_inputs = [
            input_ids,
            max_length,
            min_length,
            repetition_penalty,
        ]

    if args.vocab_mask:
        vocab_mask = onnx.helper.make_tensor_value_info("vocab_mask", TensorProto.INT32, [vocab_size])
        graph_inputs.append(vocab_mask)

    if args.prefix_vocab_mask:
        prefix_vocab_mask = onnx.helper.make_tensor_value_info(
            "prefix_vocab_mask", TensorProto.INT32, ["batch_size", vocab_size]
        )
        graph_inputs.append(prefix_vocab_mask)

    if args.custom_attention_mask:
        attention_mask = onnx.helper.make_tensor_value_info(
            "attention_mask", TensorProto.INT32, ["batch_size", "sequence_length"]
        )
        graph_inputs.append(attention_mask)

    if args.custom and args.presence_mask:
        presence_mask = onnx.helper.make_tensor_value_info(
            "presence_mask", TensorProto.INT32, ["batch_size", vocab_size]
        )
        graph_inputs.append(presence_mask)

    if is_sampling and args.seed:
        seed = onnx.helper.make_tensor_value_info("seed", TensorProto.INT32, [1])
        graph_inputs.append(seed)

    # graph outputs
    sequences = None
    if is_beamsearch:
        sequences = onnx.helper.make_tensor_value_info(
            "sequences",
            TensorProto.INT32,
            ["batch_size", "num_return_sequences", "max_length"],
        )
    elif is_greedysearch or is_sampling:
        sequences = onnx.helper.make_tensor_value_info(
            "sequences",
            TensorProto.INT32,
            ["batch_size", "max_length"],
        )

    graph_outputs = [sequences]

    if args.output_sequences_scores:
        sequences_scores = onnx.helper.make_tensor_value_info(
            "sequences_scores",
            TensorProto.FLOAT,
            ["batch_size", "num_return_sequences"],
        )
        graph_outputs.append(sequences_scores)

    if args.output_token_scores:
        scores = onnx.helper.make_tensor_value_info(
            "scores",
            TensorProto.FLOAT,
            ["max_length - sequence_length", "batch_size", "num_beams", vocab_size],
        )
        graph_outputs.append(scores)

    new_graph = onnx.helper.make_graph(
        [node],
        (f"{args.model_type} beam search" if not is_greedysearch else f"{args.model_type} greedy search"),
        graph_inputs,
        graph_outputs,
        initializers,
    )

    # Create the model
    new_model = onnx.helper.make_model(
        new_graph,
        producer_name="onnxruntime.transformers",
        opset_imports=decoder_model.opset_import,
    )

    # TODO(tianleiwu): move shared initializers from T5 encoder and decoder subgraphs to parent graph to save memory.
    if args.use_external_data_format:
        from packaging import version  # noqa: PLC0415

        if version.parse(onnx.__version__) < version.parse("1.12.0"):
            logger.warning("Require onnx >= 1.12 to save large (>2GB) model!")

        OnnxModel.save(
            new_model,
            args.output,
            save_as_external_data=True,
            all_tensors_to_one_file=True,
        )
    else:
        onnx.save(new_model, args.output)
    logger.info(f"model save to {args.output}")


def test_torch_performance(
    args: argparse.Namespace,
    model: GPT2LMHeadModel | T5ForConditionalGeneration,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor,
    eos_token_id: int,
    pad_token_id: int,
    bad_words_ids: list[list[int]],
) -> dict[str, Any]:
    """Test PyTorch performance of text generation.

    Args:
        args (argparse.Namespace): arguments parsed from command line
        model (Union[GPT2LMHeadModel, T5ForConditionalGeneration]): PyTorch model
        input_ids (torch.Tensor): input_ids
        attention_mask (torch.Tensor): Attention mask
        eos_token_id (int): EOS token ID
        pad_token_id (int): Padding token ID
        bad_words_ids (List[List[int]]): Words shall not be generated.

    Raises:
        RuntimeError: PyTorch with CUDA is not available for --use_gpu

    Returns:
        Dict[str, Any]: A dictionary with string with metric name, and value can be integer or string.
    """
    if args.use_gpu and not torch.cuda.is_available():
        raise RuntimeError("Please install PyTorch with Cuda for testing gpu performance.")

    if args.precision == Precision.FLOAT16.value:
        model.half()

    device = torch.device("cuda:0" if args.use_gpu else "cpu")
    model.to(device)

    torch.set_grad_enabled(False)
    input_ids = input_ids.to(device)
    attention_mask = attention_mask.to(device)

    torch_latency = []
    for _ in range(args.total_runs):
        start = time.time()
        _ = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=args.max_length,
            min_length=args.min_length,
            num_beams=args.num_beams,
            early_stopping=args.early_stopping,
            no_repeat_ngram_size=args.no_repeat_ngram_size,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            num_return_sequences=args.num_return_sequences,
            length_penalty=args.length_penalty,
            repetition_penalty=args.repetition_penalty,
            bad_words_ids=bad_words_ids if bad_words_ids else None,
            return_dict_in_generate=True,
            output_scores=args.output_sequences_scores or args.output_token_scores,
        )
        torch_latency.append(time.time() - start)
    batch_size = input_ids.shape[0]
    from benchmark_helper import get_latency_result  # noqa: PLC0415

    return get_latency_result(torch_latency, batch_size)


def create_attention_mask(input_ids, pad_token_id):
    attention_mask = np.ones(input_ids.shape, dtype=np.int32)
    for i in range(input_ids.shape[0]):
        abs_pos = 0
        for j in range(input_ids.shape[1]):
            if input_ids[i][j] == pad_token_id and abs_pos == 0:
                attention_mask[i][j] = 0
            else:
                abs_pos += 1
    return attention_mask


def test_gpt_model(
    args: argparse.Namespace,
    sentences: list[str] | None = None,
    is_greedy: bool = False,
):
    """Test GPT-2 model

    Args:
        args (argparse.Namespace): arguments parsed from command line
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    """
    assert args.model_type == "gpt2"

    tokenizer = GPT2Tokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
    tokenizer.padding_side = "left"
    tokenizer.pad_token = tokenizer.eos_token

    model = GPT2LMHeadModel.from_pretrained(
        args.model_name_or_path,
        cache_dir=args.cache_dir,
        pad_token_id=tokenizer.eos_token_id,
    )

    # Use different length sentences to test batching
    if sentences is None:
        sentences = [
            "The product is released",
            "I enjoy walking in the park",
            "Test best way to invest",
        ]

    inputs = tokenizer(sentences, return_tensors="pt", padding=True)
    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]

    bad_words = "walk in park"
    bad_words_ids = tokenizer.encode(bad_words, add_prefix_space=True)
    bad_words_ids = [[word_id] for word_id in bad_words_ids]  # Convert to list of list
    if args.vocab_mask:
        logger.debug("bad_words_ids", bad_words_ids)  # noqa: PLE1205
    else:
        bad_words_ids = []

    config = model.config
    eos_token_id = config.eos_token_id
    pad_token_id = config.eos_token_id
    vocab_size = config.vocab_size

    torch_decoded_sequences = []
    beam_outputs = None
    if not args.disable_parity:
        print("-" * 50)
        print("Test PyTorch model and beam search with huggingface transformers...")
        beam_outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=args.max_length,
            min_length=args.min_length,
            num_beams=args.num_beams,
            early_stopping=args.early_stopping,
            no_repeat_ngram_size=args.no_repeat_ngram_size,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            num_return_sequences=args.num_return_sequences,
            length_penalty=args.length_penalty,
            repetition_penalty=args.repetition_penalty,
            bad_words_ids=bad_words_ids if bad_words_ids else None,
            return_dict_in_generate=True,
            output_scores=args.output_sequences_scores or args.output_token_scores,
        )
        print("input_ids", input_ids)
        print("huggingface transformers outputs:")
        print("sequences", beam_outputs.sequences)
        if args.output_sequences_scores:
            print("sequences_scores", beam_outputs.sequences_scores)
        if args.output_token_scores:
            print("scores", beam_outputs.scores)
        for i, sequence in enumerate(beam_outputs.sequences):
            decoded_sequence = tokenizer.decode(sequence, skip_special_tokens=True)
            torch_decoded_sequences.append(decoded_sequence)
            print(f"{i}: {decoded_sequence}")

    print("-" * 50)
    print("Testing beam search with onnxruntime...")

    if is_greedy:
        inputs = {
            "input_ids": input_ids.cpu().numpy().astype(np.int32),
            "max_length": np.array([args.max_length], dtype=np.int32),
            "min_length": np.array([args.min_length], dtype=np.int32),
            "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32),
        }
    else:
        inputs = {
            "input_ids": input_ids.cpu().numpy().astype(np.int32),
            "max_length": np.array([args.max_length], dtype=np.int32),
            "min_length": np.array([args.min_length], dtype=np.int32),
            "num_beams": np.array([args.num_beams], dtype=np.int32),
            "num_return_sequences": np.array([args.num_return_sequences], dtype=np.int32),
            "length_penalty": np.array([args.length_penalty], dtype=np.float32),
            "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32),
        }

    if args.vocab_mask:
        vocab_mask = np.ones((vocab_size), dtype=np.int32)
        if args.vocab_mask:
            for bad_word_id in bad_words_ids:
                vocab_mask[bad_word_id] = 0
        inputs["vocab_mask"] = vocab_mask

    if args.custom_attention_mask:
        inputs["attention_mask"] = create_attention_mask(input_ids, pad_token_id)

    batch_size = input_ids.shape[0]
    if args.prefix_vocab_mask:
        logger.info("Use prefix vocab mask with all ones in ORT, but no corresponding setting for Torch model.")
        prefix_vocab_mask = np.ones((batch_size, vocab_size), dtype=np.int32)
        inputs["prefix_vocab_mask"] = prefix_vocab_mask

    if args.save_test_data:
        test_data_dir = Path(args.output).parent.as_posix()
        logger.debug("test_data_dir", test_data_dir)  # noqa: PLE1205
        from bert_test_data import output_test_data  # noqa: PLC0415

        logger.info(f"Saving test_data to {test_data_dir}/test_data_set_* ...")

        all_inputs = [inputs]
        for i, inputs in enumerate(all_inputs):
            dir = os.path.join(test_data_dir, "test_data_set_" + str(i))
            output_test_data(dir, inputs)

    logger.debug("ORT inputs", inputs)  # noqa: PLE1205

    if args.disable_perf_test:
        return

    logger.debug("Creating ort session......")
    ort_session = create_ort_session(args.output, args.use_gpu, args.use_sln_strict_mode)

    logger.debug("Run ort session......")
    result = ort_session.run(None, inputs)

    # Test performance
    latency = []
    for _ in range(args.total_runs):
        start = time.time()
        _ = ort_session.run(None, inputs)
        latency.append(time.time() - start)

    from benchmark_helper import get_latency_result  # noqa: PLC0415

    batch_size = input_ids.shape[0]
    output = get_latency_result(latency, batch_size)

    print("ORT outputs:")
    sequences = result[0]
    print("sequences", sequences)
    if args.output_sequences_scores:
        print("sequences_scores", result[1])
    if args.output_token_scores:
        print("scores", result[2])

    if is_greedy:
        (batch_size, max_length) = sequences.shape
        ort_decoded_sequences = []
        for i in range(batch_size):
            decoded_sequence = tokenizer.decode(sequences[i], skip_special_tokens=True)
            ort_decoded_sequences.append(decoded_sequence)
            print(f"batch {i} sequence: {decoded_sequence}")
    else:
        (batch_size, num_sequences, max_length) = sequences.shape
        ort_decoded_sequences = []
        for i in range(batch_size):
            for j in range(num_sequences):
                decoded_sequence = tokenizer.decode(sequences[i][j], skip_special_tokens=True)
                ort_decoded_sequences.append(decoded_sequence)
                print(f"batch {i} sequence {j}: {decoded_sequence}")

    if beam_outputs:
        torch_sequences = beam_outputs.sequences.reshape(batch_size, args.num_return_sequences, -1)
        ort_sequences = torch.LongTensor(sequences)
        print("-" * 50)
        print("Torch Sequences:")
        print(torch_sequences)
        print(torch_decoded_sequences)
        print("-" * 50)
        print("ORT Sequences:")
        print(ort_sequences)
        print(ort_decoded_sequences)
        print("-" * 50)
        # Compare the generated text instead of word IDs since ORT pads to max sequence length but Torch not.
        is_same = torch_decoded_sequences == ort_decoded_sequences
        print("Torch and ORT result is", "same" if is_same else "different")
        output["parity"] = is_same

    if args.torch_performance:
        torch_latency_output = test_torch_performance(
            args,
            model,
            input_ids,
            attention_mask,
            eos_token_id,
            pad_token_id,
            bad_words_ids,
        )
        print("Torch Latency", torch_latency_output)

    print("ORT", output)

    return output


def test_t5_model(args: argparse.Namespace, sentences: list[str] | None = None):
    """Test T5 or MT5 model

    Args:
        args (argparse.Namespace): arguments parsed from command line
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    """
    assert args.model_type in ["t5", "mt5"]

    if args.prefix_vocab_mask:
        logger.debug("Skipping parity test as prefix vocab mask is not implemented by Hugging Face")
        return None

    tokenizer = T5Tokenizer.from_pretrained(args.model_name_or_path, cache_dir=args.cache_dir)
    tokenizer.padding_side = "left"

    if args.model_type == "t5":
        model = T5ForConditionalGeneration.from_pretrained(
            args.model_name_or_path,
            cache_dir=args.cache_dir,
        )
    else:
        model = MT5ForConditionalGeneration.from_pretrained(
            args.model_name_or_path,
            cache_dir=args.cache_dir,
        )

    # Use different length sentences to test batching
    if sentences is None:
        sentences = [
            "translate English to French: The product is released",
            "summarize: research continues to show that pets bring real health benefits to their owners. Having a dog around can lead to lower levels of stress for both adults and kids.",
            # "summarize: I enjoy walking in the park. It makes my mind feel calm and refreshed. "
            # + "I enjoy looking at the trees, flowers, and wildlife around me, and listening to sound from natural.",
        ]

    inputs = tokenizer(sentences, return_tensors="pt", padding=True)
    input_ids = inputs["input_ids"]
    attention_mask = inputs["attention_mask"]

    bad_words = "walk in park"
    bad_words_ids = tokenizer.encode(bad_words)[:-1]  # exclude the last token (EOS)
    bad_words_ids = [[word_id] for word_id in bad_words_ids]  # Convert to list of list
    if args.vocab_mask:
        logger.debug("bad_words_ids", bad_words_ids)  # noqa: PLE1205
    else:
        bad_words_ids = []

    config = model.config
    eos_token_id = config.eos_token_id
    pad_token_id = config.pad_token_id
    vocab_size = config.vocab_size
    logger.debug(f"eos_token_id:{eos_token_id}, pad_token_id:{pad_token_id}, vocab_size:{vocab_size}")

    torch_decoded_sequences = []
    if not args.disable_parity:
        print("-" * 50)
        print("Test PyTorch model and beam search with huggingface transformers...")
        beam_outputs = model.generate(
            input_ids=input_ids,
            attention_mask=attention_mask,
            max_length=args.max_length,
            min_length=args.min_length,
            num_beams=args.num_beams,
            early_stopping=args.early_stopping,
            no_repeat_ngram_size=args.no_repeat_ngram_size,
            eos_token_id=eos_token_id,
            pad_token_id=pad_token_id,
            num_return_sequences=args.num_return_sequences,
            length_penalty=args.length_penalty,
            repetition_penalty=args.repetition_penalty,
            bad_words_ids=bad_words_ids if bad_words_ids else None,
            return_dict_in_generate=True,
            output_scores=args.output_sequences_scores or args.output_token_scores,
        )

        print("input_ids", input_ids)
        print("huggingface transformers outputs:")
        print("sequences", beam_outputs.sequences)
        if args.output_sequences_scores:
            print("sequences_scores", beam_outputs.sequences_scores)
        if args.output_token_scores:
            print("scores", beam_outputs.scores)
        for i, sequence in enumerate(beam_outputs.sequences):
            decoded_sequence = tokenizer.decode(sequence, skip_special_tokens=True)
            torch_decoded_sequences.append(decoded_sequence)
            print(f"{i}: {decoded_sequence}")

    print("-" * 50)
    print("Testing beam search with onnxruntime...")

    vocab_mask = np.ones((vocab_size), dtype=np.int32)
    if args.vocab_mask:
        for bad_word_id in bad_words_ids:
            vocab_mask[bad_word_id] = 0

    inputs = {
        "input_ids": input_ids.cpu().numpy().astype(np.int32),
        "max_length": np.array([args.max_length], dtype=np.int32),
        "min_length": np.array([args.min_length], dtype=np.int32),
        "num_beams": np.array([args.num_beams], dtype=np.int32),
        "num_return_sequences": np.array([args.num_return_sequences], dtype=np.int32),
        "length_penalty": np.array([args.length_penalty], dtype=np.float32),
        "repetition_penalty": np.array([args.repetition_penalty], dtype=np.float32),
    }

    if args.vocab_mask:
        inputs["vocab_mask"] = vocab_mask

    if args.custom_attention_mask:
        inputs["attention_mask"] = create_attention_mask(input_ids, pad_token_id)

    if args.save_test_data:
        test_data_dir = Path(args.output).parent.as_posix()
        logger.debug("test_data_dir", test_data_dir)  # noqa: PLE1205
        from bert_test_data import output_test_data  # noqa: PLC0415

        all_inputs = [inputs]
        for i, inputs in enumerate(all_inputs):
            dir = os.path.join(test_data_dir, "test_data_set_" + str(i))
            output_test_data(dir, inputs)

    logger.debug("ORT inputs", inputs)  # noqa: PLE1205

    ort_session = create_ort_session(args.output, args.use_gpu, args.use_sln_strict_mode)

    # Test performance
    latency = []
    for _ in range(args.total_runs):
        start = time.time()
        result = ort_session.run(None, inputs)
        latency.append(time.time() - start)
    batch_size = input_ids.shape[0]
    from benchmark_helper import get_latency_result  # noqa: PLC0415

    output = get_latency_result(latency, batch_size)

    print("ORT outputs:")
    sequences = result[0]
    print("sequences", sequences)
    if args.output_sequences_scores:
        print("sequences_scores", result[1])
    if args.output_token_scores:
        print("scores", result[2])

    (batch_size, num_sequences, max_length) = sequences.shape
    ort_decoded_sequences = []
    for i in range(batch_size):
        for j in range(num_sequences):
            decoded_sequence = tokenizer.decode(sequences[i][j], skip_special_tokens=True)
            ort_decoded_sequences.append(decoded_sequence)
            print(f"batch {i} sequence {j}: {decoded_sequence}")

    if not args.disable_parity:
        torch_sequences = beam_outputs.sequences.reshape(batch_size, args.num_return_sequences, -1)
        ort_sequences = torch.LongTensor(sequences)
        print("-" * 50)
        print("Torch Sequences:")
        print(torch_sequences)
        print(torch_decoded_sequences)
        print("-" * 50)
        print("ORT Sequences:")
        print(ort_sequences)
        print(ort_decoded_sequences)
        print("-" * 50)
        # Compare the generated text instead of word IDs since ORT pads to max sequence length but Torch not.
        is_same = torch_decoded_sequences == ort_decoded_sequences
        print("Torch and ORT result is ", "same" if is_same else "different")
        output["parity"] = is_same

    if args.torch_performance:
        torch_latency_output = test_torch_performance(
            args,
            model,
            input_ids,
            attention_mask,
            eos_token_id,
            pad_token_id,
            bad_words_ids,
        )
        print("Torch Latency", torch_latency_output)

    print("ORT", output)
    return output


def main(argv: list[str] | None = None, sentences: list[str] | None = None):
    """Main entry function

    Args:
        argv (Optional[List[str]], optional): _description_. Defaults to None.
        sentences (Optional[List[str]], optional): input text. Defaults to None.

    Raises:
        ValueError: Path does not exist: --encoder_decoder_init_onnx
        ValueError: Path does not exist: --decoder_onnx
        ValueError: --decoder_onnx and --encoder_decoder_init_onnx are not used together for T5

    Returns:
        Union[Dict[str, Any], None]: A dictionary with string with metric name, and value can be integer or string.
    """

    args = parse_arguments(argv)
    setup_logger(args.verbose)

    if args.model_type in ["t5", "mt5"]:
        if args.encoder_decoder_init_onnx and not os.path.exists(args.encoder_decoder_init_onnx):
            raise ValueError(f"Path does not exist: --encoder_decoder_init_onnx {args.encoder_decoder_init_onnx}")
        if args.decoder_onnx and not os.path.exists(args.decoder_onnx):
            raise ValueError(f"Path does not exist: --decoder_onnx {args.decoder_onnx}")
        if (args.encoder_decoder_init_onnx and not args.decoder_onnx) or (
            args.decoder_onnx and not args.encoder_decoder_init_onnx
        ):
            raise ValueError("--decoder_onnx shall use together with --encoder_decoder_init_onnx")

    is_greedy = args.num_beams == 1 and args.num_return_sequences == 1

    if args.model_type == "gpt2" and is_greedy:
        if args.top_p > 0.0 and args.top_p < 1.0:
            convert_generation_model(args, GenerationType.SAMPLING)
            logger.info(
                "The test for gpt2_sampling onnx model is limited to non-custom model with small top_p(e.g <=0.01) value. The result should be the same as gpt2 greedy search."
            )
            if args.top_p > 0.01 or args.custom or args.seed:
                return
        else:
            convert_generation_model(args, GenerationType.GREEDYSEARCH)
    else:
        convert_generation_model(args)

    logger.info("start testing model...")
    if args.model_type in ["t5", "mt5"]:
        result = test_t5_model(args, sentences=sentences)
    else:
        result = test_gpt_model(args, sentences=sentences, is_greedy=is_greedy)

    if result:
        if args.use_external_data_format:
            logger.info(f"Output files: {args.output}, {args.output}.data")
        else:
            logger.info(f"Output file: {args.output}")

    return result


if __name__ == "__main__":
    main()