onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/onnx_model_bert_tf.py

@@ -0,0 +1,588 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+import logging
+
+import numpy as np
+import onnx
+from onnx import TensorProto, helper, numpy_helper
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class BertOnnxModelTF(BertOnnxModel):
+    def __init__(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
+
+    def remove_identity(self):
+        nodes_to_remove = []
+        for node in self.nodes():
+            if node.op_type == "Identity":
+                if not self.find_graph_output(node.output[0]):
+                    self.replace_input_of_all_nodes(node.output[0], node.input[0])
+                    nodes_to_remove.append(node)
+        self.remove_nodes(nodes_to_remove)
+        logger.info(f"Removed Identity count: {len(nodes_to_remove)}")
+
+    def match_mask_path(self, add_or_sub_before_softmax):
+        mask_nodes = self.match_parent_path(
+            add_or_sub_before_softmax,
+            ["Mul", "Sub", "Reshape", "Cast"],
+            [1, None, 1, 0],
+        )
+        if mask_nodes is not None:
+            return mask_nodes
+
+        mask_nodes = self.match_parent_path(
+            add_or_sub_before_softmax,
+            ["Mul", "Sub", "Cast", "Slice", "Unsqueeze"],
+            [1, 0, 1, 0, 0],
+        )
+        if mask_nodes is not None:
+            return mask_nodes
+
+        mask_nodes = self.match_parent_path(
+            add_or_sub_before_softmax,
+            ["Mul", "Sub", "Cast", "Unsqueeze", "Unsqueeze"],
+            [1, None, 1, 0, 0],
+        )
+
+        return mask_nodes
+
+    def get_2d_initializers_from_parent_subgraphs(self, current_node):
+        """
+        Find initializers that is 2D. Returns a dictionary with name as key and shape as value.
+        """
+        parent_nodes = self.get_parent_subgraph_nodes(current_node, [])
+        initializers = {}
+        for node in parent_nodes:
+            for input in node.input:
+                initializer = self.get_initializer(input)
+                if initializer:
+                    temp = numpy_helper.to_array(initializer)
+                    if len(temp.shape) == 2:
+                        initializers[initializer.name] = temp.shape
+
+        return initializers
+
+    def find_segment_ids(self, segment_embedding, input_ids):
+        input_name_to_nodes = self.input_name_to_nodes()
+        if segment_embedding not in input_name_to_nodes:
+            return None
+
+        nodes = input_name_to_nodes[segment_embedding]
+        if len(nodes) != 1:
+            return None
+
+        graph_inputs = self.get_graph_inputs(nodes[0], recursive=True)
+        if len(graph_inputs) > 1:
+            print("Found multiple candidates of segment_ids", graph_inputs)
+            return None
+        # Find segment ids in graph inputs. The segment id input must not be the same as input_ids.
+        if len(graph_inputs) == 1 and graph_inputs[0] != input_ids:
+            return graph_inputs[0]
+
+        # If the segment id candidate is the same as the input_ids, try to assign alternative segment ids and simplify the graph if needed.
+        segment_ids = nodes[0].input[1]
+        _, segment_id_path, _ = self.match_parent_paths(
+            nodes[0],
+            [
+                (
+                    ["ConstantOfShape", "Cast", "Concat", "Slice", "Cast", "Shape"],
+                    [1, 0, 0, 0, 0, 0],
+                ),
+                (
+                    [
+                        "ConstantOfShape",
+                        "Cast",
+                        "Concat",
+                        "Unsqueeze",
+                        "Squeeze",
+                        "Slice",
+                        "Cast",
+                        "Shape",
+                    ],
+                    [1, 0, 0, 0, 0, 0, 0, 0],
+                ),
+            ],
+            None,
+        )
+
+        if segment_id_path and input_ids and input_ids == segment_id_path[-1].input[0]:
+            logger.debug("Simplify semgent id path...")
+            constantofshape_node = segment_id_path[0]
+            graph_name = self.get_graph_by_node(constantofshape_node).name
+            self.add_node(
+                helper.make_node("Shape", inputs=[input_ids], outputs=["input_shape"]),
+                graph_name,
+            )
+            constantofshape_value = helper.get_attribute_value(constantofshape_node.attribute[0])
+            self.add_node(
+                helper.make_node(
+                    "ConstantOfShape",
+                    inputs=["input_shape"],
+                    outputs=["zeros_for_input_shape"],
+                    value=constantofshape_value,
+                ),
+                graph_name,
+            )
+            segment_ids = "zeros_for_input_shape"
+        return segment_ids
+
+    def find_input_ids(self, word_embedding):
+        input_name_to_nodes = self.input_name_to_nodes()
+        if word_embedding not in input_name_to_nodes:
+            return None
+
+        nodes = input_name_to_nodes[word_embedding]
+        if len(nodes) != 1:
+            return None
+
+        graph_inputs = self.get_graph_inputs(nodes[0], recursive=True)
+        if len(graph_inputs) == 1:
+            return graph_inputs[0]
+
+        print("Found multiple candidates of input_ids", graph_inputs)
+        return None
+
+    def find_mask_input(self, excluded_graph_inputs):
+        for node in self.nodes():
+            if node.op_type == "Softmax":
+                mask_path = self.match_parent_path(
+                    node,
+                    ["Add", "Mul", "Sub", "Cast", "Slice", "Unsqueeze"],
+                    [0, 1, None, 1, 0, 0],
+                )
+                if mask_path is None:
+                    continue
+                (
+                    add_node,
+                    mul_node,
+                    sub_node,
+                    cast_node,
+                    slice_node,
+                    unsqueeze_node,
+                ) = mask_path
+                if self.has_constant_input(mul_node, -10000) and self.has_constant_input(sub_node, 1):
+                    graph_inputs = self.get_graph_inputs(sub_node, recursive=True)
+                    inputs = [input for input in graph_inputs if input not in excluded_graph_inputs]
+                    if len(inputs) > 1:
+                        print("Found multiple candidates of mask input", inputs)
+                        return None
+                    if len(inputs) == 1:
+                        return inputs[0]
+                    # Duplicated input found. Try to simplify the graph.
+                    path_to_be_simplified = self.match_parent_path(
+                        mask_path[-1],
+                        [
+                            "ConstantOfShape",
+                            "Cast",
+                            "Concat",
+                            "Unsqueeze",
+                            "Squeeze",
+                            "Slice",
+                            "Cast",
+                            "Shape",
+                        ],
+                        [0, 0, 0, 0, 0, 0, 0, 0],
+                    )
+                    duplicated_inputs = [input for input in graph_inputs if input in excluded_graph_inputs]
+                    # Simplify graph for dynamic axes.
+                    if (
+                        path_to_be_simplified
+                        and duplicated_inputs
+                        and len(duplicated_inputs) == 1
+                        and duplicated_inputs[0] == path_to_be_simplified[-1].input[0]
+                    ):
+                        logger.debug("Simplify semgent id path...")
+                        constantofshape_node = path_to_be_simplified[0]
+                        constantofshape_value = helper.get_attribute_value(constantofshape_node.attribute[0])
+                        graph_name = self.get_graph_by_node(constantofshape_node).name
+                        self.add_node(
+                            helper.make_node(
+                                "Shape",
+                                inputs=[duplicated_inputs[0]],
+                                outputs=["input_shape_for_mask"],
+                            ),
+                            graph_name,
+                        )
+                        self.add_node(
+                            helper.make_node(
+                                "ConstantOfShape",
+                                inputs=["input_shape_for_mask"],
+                                outputs=[unsqueeze_node.input[0]],
+                                value=constantofshape_value,
+                            ),
+                            graph_name,
+                        )
+                    return unsqueeze_node.input[0]
+        return None
+
+    def create_embedding_subgraph(self, normalize_node, word_embedding, segment_embedding, position_embedding):
+        input_ids = self.find_input_ids(word_embedding)
+        if input_ids is None:
+            logger.info("Failed to find input_ids. Cannot fuse embedding layer.")
+            return False
+
+        segment_ids = self.find_segment_ids(segment_embedding, input_ids)
+        if segment_ids is None:
+            logger.info("Failed to find segment_ids. Cannot fuse embedding layer.")
+            return False
+
+        mask_input = self.find_mask_input([segment_ids, input_ids])
+        if mask_input is None:
+            logger.info("Failed to find input_mask. Cannot fuse embedding layer.")
+            return False
+
+        self.bert_inputs = [input_ids, segment_ids, mask_input]
+
+        mask_index = self.create_node_name("mask_index")
+        self.attention_mask.set_mask_indice(mask_input, mask_index)
+
+        if self.find_graph_input(input_ids).type.tensor_type.elem_type != TensorProto.INT32:
+            casted, input_ids = self.utils.cast_graph_input_to_int32(input_ids)
+
+        if self.find_graph_input(segment_ids):
+            casted, segment_ids = self.utils.cast_graph_input_to_int32(segment_ids)
+        else:
+            segment_ids, segment_id_cast_node = self.utils.cast_input_to_int32(segment_ids)
+
+        if self.find_graph_input(mask_input):
+            casted, mask_input = self.utils.cast_graph_input_to_int32(mask_input)
+        else:
+            mask_input, mask_input_cast_node = self.utils.cast_input_to_int32(mask_input)
+
+        embed_output = self.create_node_name("embed_output")
+        embed_node = onnx.helper.make_node(
+            "EmbedLayerNormalization",
+            inputs=[
+                input_ids,
+                segment_ids,
+                word_embedding,
+                position_embedding,
+                segment_embedding,
+                normalize_node.input[1],  # gamma
+                normalize_node.input[2],  # beta
+                mask_input,
+            ],
+            outputs=[embed_output, mask_index],
+            name="EmbedLayer",
+        )
+        embed_node.domain = "com.microsoft"
+        self.replace_input_of_all_nodes(normalize_node.output[0], embed_output)
+        self.add_node(embed_node, self.get_graph_by_node(normalize_node).name)
+
+    def process_embedding(self):
+        """
+        Automatically detect word, segment and position embeddings.
+        """
+        logger.info("start processing embedding layer...")
+        output_name_to_node = self.output_name_to_node()
+
+        layer_norm_nodes = self.get_nodes_by_op_type("LayerNormalization")
+        for layer_norm_node in layer_norm_nodes:
+            pos_embed_path = self.match_parent_path(
+                layer_norm_node,
+                ["Add", "Reshape", "Slice"],
+                [0, 1, 0],
+                output_name_to_node,
+            )
+            if pos_embed_path is None:
+                continue
+
+            add_node, reshape_node, slice_node = pos_embed_path
+            initializer = self.get_initializer(slice_node.input[0])
+            if initializer is None:
+                continue
+
+            temp = numpy_helper.to_array(initializer)
+            if len(temp.shape) == 2:
+                logger.info(f"Found position embedding. name:{initializer.name}, shape:{temp.shape}")
+                position_embedding = initializer.name
+            else:
+                logger.info(f"Failed to find position embedding. name:{initializer.name}, shape:{temp.shape}")
+                return
+
+            first_parent = self.get_parent(add_node, 0, output_name_to_node)
+            if first_parent is not None and first_parent.op_type == "Add":
+                embeddings = self.get_2d_initializers_from_parent_subgraphs(first_parent)
+                if len(embeddings) != 2:
+                    logger.warning(
+                        f"Failed to find two embeddings (word and segment) from Add node. Found {embeddings}"
+                    )
+                    return
+
+                word_embedding = None
+                segment_embedding = None
+                for name, shape in embeddings.items():
+                    if shape[0] == 2:
+                        segment_embedding = name
+                        logger.info(f"Found segment embedding. name:{name}, shape:{shape}")
+                    else:
+                        word_embedding = name
+                        logger.info(f"Found words embedding. name:{name}, shape:{shape}")
+
+                if word_embedding is None or segment_embedding is None:
+                    logger.info("Failed to find both word and segment embedding")
+                    return
+
+                logger.info("Create Embedding node")
+                self.create_embedding_subgraph(
+                    layer_norm_node,
+                    word_embedding,
+                    segment_embedding,
+                    position_embedding,
+                )
+                # Prune graph to remove those original embedding nodes.
+                self.prune_graph()
+                break
+
+    def check_attention_input(self, matmul_q, matmul_k, matmul_v, parent, output_name_to_node):
+        for x in [matmul_q, matmul_k, matmul_v]:
+            root_input = x.input[0]
+            root_node = output_name_to_node[root_input]
+            if root_node == parent:
+                continue
+            logger.debug(f"Check attention input failed:{root_input}, {parent.output[0]}")
+            return False
+
+        return True
+
+    def fuse_attention(self):
+        output_name_to_node = self.output_name_to_node()
+
+        nodes_to_remove = []
+        attention_count = 0
+
+        start_nodes = []
+        skip_layer_norm_nodes = self.get_nodes_by_op_type("SkipLayerNormalization")
+        layer_norm_nodes = self.get_nodes_by_op_type("LayerNormalization")
+        # Sometimes we can not fuse skiplayernormalization since the add before layernorm has an output that used by nodes outside skiplayernorm
+        # Conceptually we treat add before layernorm as skiplayernorm node since they share the same pattern
+        start_nodes.extend(skip_layer_norm_nodes)
+        start_nodes.extend(layer_norm_nodes)
+
+        for normalize_node in start_nodes:
+            graph_name = self.get_graph_by_node(normalize_node).name
+            # SkipLayerNormalization has two inputs, and one of them is the root input for attention.
+            if normalize_node.op_type == "LayerNormalization":
+                add_before_layernorm = self.match_parent(normalize_node, "Add", 0)
+                if add_before_layernorm is not None:
+                    normalize_node = add_before_layernorm  # noqa: PLW2901
+                else:
+                    continue
+            parent = self.get_parent(normalize_node, 1)
+            if parent is None or parent.op_type not in [
+                "SkipLayerNormalization",
+                "LayerNormalization",
+                "Reshape",
+            ]:
+                parent = self.get_parent(normalize_node, 0)
+                if parent is None or parent.op_type not in [
+                    "SkipLayerNormalization",
+                    "LayerNormalization",
+                    "Reshape",
+                ]:
+                    logger.debug("Failed to match parent of normalize_node")
+                    continue
+
+            qkv_nodes = self.match_parent_path(
+                normalize_node,
+                ["Add", "MatMul", "Reshape", "Transpose", "MatMul"],
+                [0, 0, 0, 0, 0],
+            )
+            if qkv_nodes is None:
+                qkv_nodes = self.match_parent_path(
+                    normalize_node,
+                    ["MatMul", "Reshape", "Transpose", "MatMul"],
+                    [1, 0, 0, 0],
+                )
+                if qkv_nodes is None:
+                    qkv_nodes = self.match_parent_path(normalize_node, ["Add", "Einsum", "Einsum"], [0, 0, 0])
+                    if qkv_nodes is None:
+                        logger.debug("Failed to match qkv nodes")
+                        continue
+
+            matmul_qkv = qkv_nodes[-1]
+            v_nodes = self.match_parent_path(matmul_qkv, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0])
+            if v_nodes is None:
+                v_nodes = self.match_parent_path(matmul_qkv, ["Add", "Einsum"], [1, 0])
+                if v_nodes is None:
+                    logger.debug("Failed to match v path")
+                    continue
+
+            add_v = v_nodes[-2]
+            matmul_v = v_nodes[-1]
+            qk_nodes = self.match_parent_path(matmul_qkv, ["Softmax", "Add", "Mul", "MatMul"], [0, 0, 0, 0])
+            if qk_nodes is None:
+                qk_nodes = self.match_parent_path(matmul_qkv, ["Softmax", "Add", "Einsum"], [0, 0, 0])
+                if qk_nodes is None:
+                    logger.debug("Failed to match qk_paths")
+                    continue
+            matmul_qk = qk_nodes[-1]
+
+            q_nodes = self.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [0, 0, 0, 0])
+            if q_nodes is None:
+                q_nodes = self.match_parent_path(matmul_qk, ["Add", "Einsum"], [0, 0])
+                if q_nodes is None:
+                    logger.debug("Failed to match q path")
+                    continue
+
+            add_q = q_nodes[-2]
+            matmul_q = q_nodes[-1]
+
+            k_nodes = self.match_parent_path(matmul_qk, ["Transpose", "Reshape", "Add", "MatMul"], [1, 0, 0, 0])
+            if k_nodes is None:
+                k_nodes = self.match_parent_path(matmul_qk, ["Mul", "Add", "Einsum"], [1, 0, 0])
+                if k_nodes is None:
+                    logger.debug("Failed to match k path")
+                    continue
+            add_k = k_nodes[-2]
+            matmul_k = k_nodes[-1]
+
+            mask_nodes = self.match_mask_path(qk_nodes[1])
+
+            if mask_nodes is None:
+                logger.debug("Cannot find mask_nodes.")
+                continue
+
+            if not self.has_constant_input(mask_nodes[1], 1):
+                logger.debug("Sub node expected to have an input with constant value 1.0.")
+                continue
+
+            # add a squeeze node to convert a 3-d mask to 2-d
+            squeeze_node = self.match_parent_path(mask_nodes[-1], ["Squeeze"], [0]) or self.match_parent_path(
+                mask_nodes[-1], ["Expand"], [0]
+            )
+            squeeze_node_name = "Squeeze_3d_to_2d_mask"
+            squeeze_output_name = squeeze_node_name + "_output"
+            if squeeze_node is None and len(mask_nodes) == 5 and self.find_graph_input(mask_nodes[-1].input[0]) is None:
+                mask_input = mask_nodes[-1].input[1]
+                self.add_node(
+                    helper.make_node(
+                        "Squeeze",
+                        [mask_input],
+                        [squeeze_output_name],
+                        squeeze_node_name,
+                        axes=[1],
+                    ),
+                    graph_name,
+                )
+                mask_nodes[-1].input[0] = squeeze_output_name
+
+            is_same_root = self.check_attention_input(matmul_q, matmul_k, matmul_v, parent, output_name_to_node)
+            if is_same_root:
+                mask_index = self.attention_mask.process_mask(mask_nodes[-1].input[0])
+                logger.debug("Create an Attention node.")
+
+                # For tf models, q and v are flipped.
+                attention_node = self.attention_fusion.create_attention_node(
+                    mask_index=mask_index,
+                    q_matmul=matmul_k,
+                    k_matmul=matmul_q,
+                    v_matmul=matmul_v,
+                    q_add=add_k,
+                    k_add=add_q,
+                    v_add=add_v,
+                    num_heads=self.num_heads,
+                    hidden_size=self.hidden_size,
+                    first_input=parent.output[0],
+                    output=qkv_nodes[2].output[0],
+                )
+                if attention_node is None:
+                    continue
+
+                if qkv_nodes[1].op_type == "Einsum":
+                    # add reshape before einsum
+                    tensor = helper.make_tensor(
+                        name=qkv_nodes[1].name + "_newshape",
+                        data_type=TensorProto.INT64,
+                        dims=[4],
+                        vals=np.int64(
+                            [
+                                [
+                                    0,
+                                    0,
+                                    self.num_heads,
+                                    int(self.hidden_size / self.num_heads),
+                                ]
+                            ]
+                        ).tobytes(),
+                        raw=True,
+                    )
+                    self.add_initializer(tensor, graph_name)
+                    reshape_ = helper.make_node(
+                        "Reshape",
+                        inputs=[
+                            attention_node.output[0],
+                            qkv_nodes[1].name + "_newshape",
+                        ],
+                        outputs=[qkv_nodes[1].name + "_reshape_output"],
+                        name=qkv_nodes[1].name + "_reshape",
+                    )
+                    qkv_nodes[1].input[0] = qkv_nodes[1].name + "_reshape_output"
+                    self.add_node(reshape_, graph_name)
+                if parent.op_type == "Reshape":
+                    # Temporary work around: we require the skiplayernorm and attention op be fed with 3-d input
+                    hidden_size = numpy_helper.to_array(self.get_initializer(parent.input[1]))[1]
+                    tensor = helper.make_tensor(
+                        name=parent.name + "_modified",
+                        data_type=TensorProto.INT64,
+                        dims=[3],
+                        vals=np.int64([[1, -1, hidden_size]]).tobytes(),
+                        raw=True,
+                    )
+                    self.add_initializer(tensor, graph_name)
+                    parent.input[1] = parent.name + "_modified"
+
+                self.add_node(attention_node, graph_name)
+                attention_count += 1
+
+                nodes_to_remove.extend(qkv_nodes[2:])
+                nodes_to_remove.extend(qk_nodes)
+                nodes_to_remove.extend(q_nodes)
+                nodes_to_remove.extend(k_nodes)
+                nodes_to_remove.extend(v_nodes)
+                nodes_to_remove.extend(mask_nodes)
+            else:
+                logger.debug("Root node not matched.")
+                continue
+        self.remove_nodes(nodes_to_remove)
+        self.update_graph()
+        logger.info(f"Fused Attention count:{attention_count}")
+
+    def preprocess(self):
+        self.remove_identity()
+        self.process_embedding()
+        self.skip_reshape()
+
+    def skip_reshape(self):
+        count = 0
+        reshape_nodes = self.get_nodes_by_op_type("Reshape")
+        for reshape_node in reshape_nodes:
+            parent = self.get_parent(reshape_node, 0)
+            if parent is not None and parent.op_type == "Reshape":
+                reshape_node.input[0] = parent.input[0]
+                count += 1
+
+        if count > 0:
+            logger.info(f"Skip consequent Reshape count: {count}")
+
+    def remove_reshape_before_first_attention(self):
+        attention_nodes = self.get_nodes_by_op_type("Attention")
+        for attention_node in attention_nodes:
+            path = self.match_parent_path(attention_node, ["Reshape", "EmbedLayerNormalization"], [0, 0])
+            if path is None:
+                continue
+            logger.info("Remove Reshape before first Attention node.")
+            reshape, _ = path
+            self.replace_input_of_all_nodes(reshape.output[0], reshape.input[0])
+            self.remove_node(reshape)
+            break
+
+    def postprocess(self):
+        self.remove_reshape_before_first_attention()
+        self.prune_graph()

onnxruntime/transformers/onnx_model_clip.py

@@ -0,0 +1,42 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+from fusion_attention_clip import FusionAttentionClip
+from onnx import ModelProto
+from onnx_model_bert import BertOnnxModel
+
+logger = getLogger(__name__)
+
+
+class ClipOnnxModel(BertOnnxModel):
+    def __init__(self, model: ModelProto, num_heads: int = 0, hidden_size: int = 0):
+        super().__init__(model, num_heads=num_heads, hidden_size=hidden_size)
+        self.clip_attention_fusion = FusionAttentionClip(self, self.hidden_size, self.num_heads)
+
+    def get_fused_operator_statistics(self):
+        """
+        Returns node count of fused operators.
+        """
+        op_count = {}
+        ops = [
+            "Attention",
+            "FastGelu",
+            "Gelu",
+            "LayerNormalization",
+            "QuickGelu",
+            "BiasGelu",
+            "SkipLayerNormalization",
+        ]
+        for op in ops:
+            nodes = self.get_nodes_by_op_type(op)
+            op_count[op] = len(nodes)
+
+        logger.info(f"Optimized operators:{op_count}")
+        return op_count
+
+    def fuse_attention(self):
+        self.clip_attention_fusion.apply()

onnxruntime/transformers/onnx_model_conformer.py

@@ -0,0 +1,32 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+import logging
+
+from fusion_attention import AttentionMask
+from fusion_conformer_attention import FusionConformerAttention
+from fusion_options import FusionOptions
+from onnx_model_bert import BertOnnxModel
+
+logger = logging.getLogger(__name__)
+
+
+class ConformerOnnxModel(BertOnnxModel):
+    def __init__(self, model, num_heads, hidden_size):
+        super().__init__(model, num_heads, hidden_size)
+        self.attention_mask = AttentionMask(self)
+        self.attention_fusion = FusionConformerAttention(self, self.hidden_size, self.num_heads, self.attention_mask)
+
+    def optimize(self, options: FusionOptions | None = None, add_dynamic_axes: bool = False):
+        self.attention_fusion.use_multi_head_attention = False if options is None else options.use_multi_head_attention
+        self.attention_fusion.disable_multi_head_attention_bias = (
+            False if options is None else options.disable_multi_head_attention_bias
+        )
+        super().optimize(options, add_dynamic_axes)
+
+    def fuse_attention(self):
+        self.attention_fusion.apply()
+
+    def preprocess(self):
+        self.adjust_reshape_and_expand()