onnxruntime-directml 1.24.1__cp314-cp314-win_amd64.whl

onnxruntime/transformers/fusion_skiplayernorm.py

@@ -0,0 +1,209 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+from fusion_base import Fusion
+from fusion_utils import NumpyHelper
+from onnx import helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionSkipLayerNormalization(Fusion):
+    """
+    Fuse Add + LayerNormalization into one node: SkipLayerNormalization
+    Note: This fusion does not check the input shape of Add and LayerNormalization.
+    """
+
+    def __init__(
+        self,
+        model: OnnxModel,
+        fused_op_type: str = "SkipLayerNormalization",
+        search_op_types: str = "LayerNormalization",
+        shape_infer: bool = True,
+    ):
+        super().__init__(model, fused_op_type, search_op_types)
+        if shape_infer:
+            # Update shape inference is needed since other fusions might add new edge which does not have shape info yet.
+            self.shape_infer_helper = self.model.infer_runtime_shape({"batch_size": 4, "seq_len": 7}, update=True)
+            if self.shape_infer_helper is None:
+                # TODO(tianleiwu): support subgraph in shape inference or add broadcasting in SkipLayerNormalization op.
+                logger.warning("symbolic shape inference disabled or failed.")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        add = self.model.get_parent(node, 0, output_name_to_node)
+
+        # In some models there is input_ids->gather->add->LayerNorm and one of input of the
+        # add node is initializer with fixed shape which should not be fused into SkipLayerNorm
+        if add is None or add.op_type != "Add":
+            return
+
+        # The number of inputs of add should be 2
+        if len(add.input) != 2:
+            return
+
+        for add_input in add.input:
+            if self.model.get_initializer(add_input) is not None:
+                return
+
+        # To avoid an Add node have two children of LayerNormalization, we shall only fuse one SkipLayerNormalization
+        if add in self.nodes_to_remove:
+            return
+
+        # Root Mean Square Layer Normalization
+        simplified = node.op_type == "SimplifiedLayerNormalization"
+
+        if hasattr(self, "shape_infer_helper"):
+            if self.shape_infer_helper is not None:
+                if (
+                    self.shape_infer_helper.get_edge_shape(add.input[0])
+                    and len(self.shape_infer_helper.get_edge_shape(add.input[0])) != 3
+                ):
+                    logger.debug("skip SkipLayerNormalization fusion since shape of input %s is not 3D", add.input[0])
+                    return
+
+                # TODO(tianleiwu): support broadcasting Skip shape (1, sequence_length, hidden_size) or (sequence_length, hidden_size)
+                if not self.shape_infer_helper.compare_shape(add.input[0], add.input[1]):
+                    logger.debug(
+                        "skip SkipLayerNormalization fusion since shape of inputs (%s, %s) are not same",
+                        add.input[0],
+                        add.input[1],
+                    )
+                    return
+            else:
+                logger.debug("skip SkipLayerNormalization fusion since symbolic shape inference failed")
+                return
+
+        gather_path = self.model.match_parent_path(add, ["Gather"], [None])
+        if gather_path is not None and self.model.find_graph_input(gather_path[0].input[1]) is None:
+            if self.model.match_parent_path(gather_path[0], ["ConstantOfShape"], [1]) is None:
+                return
+
+        # This means that the residual Add before the LayerNormalization produces an output
+        # that is consumed by some other nodes or graph output other than the LayerNormalization itself
+        # We can still go ahead with the SkipLayerNormalization fusion but we need to
+        # preserve the output of Add and that needs to be produced by SkipLayerNormalization.
+        add_has_graph_output = self.model.find_graph_output(add.output[0]) is not None
+        residual_add_has_multiple_consumers = (
+            add_has_graph_output or len(self.model.get_children(add, input_name_to_nodes)) > 1
+        )
+
+        outputs_to_keep = node.output
+
+        if residual_add_has_multiple_consumers:
+            outputs_to_keep.extend([add.output[0]])
+
+        outputs = [node.output[0]]
+
+        # Skip the other optional outputs of SkipLayerNormalization before adding the Add's output
+        if residual_add_has_multiple_consumers:
+            outputs.extend(["", "", add.output[0]])
+
+        if self.model.is_safe_to_fuse_nodes([add, node], outputs_to_keep, input_name_to_nodes, output_name_to_node):
+            self.nodes_to_remove.extend([add, node])
+
+            inputs = (
+                [add.input[0], add.input[1], node.input[1], node.input[2]]
+                if not simplified
+                else [add.input[0], add.input[1], node.input[1]]
+            )
+            normalize_node = helper.make_node(
+                self.fused_op_type,
+                inputs=inputs,
+                outputs=outputs,
+                name=self.model.create_node_name(self.fused_op_type, name_prefix="SkipLayerNorm"),
+            )
+            normalize_node.domain = "com.microsoft"
+
+            # Pass attribute "epsilon" from layernorm node to SkipLayerNormalization
+            for att in node.attribute:
+                if att.name == "epsilon":
+                    normalize_node.attribute.extend([att])
+
+            # Set default epsilon if no epsilon exists from layernorm
+            if len(normalize_node.attribute) == 0:
+                normalize_node.attribute.extend([helper.make_attribute("epsilon", 1.0e-12)])
+
+            self.nodes_to_add.append(normalize_node)
+            self.node_name_to_graph_name[normalize_node.name] = self.this_graph_name
+
+
+class FusionBiasSkipLayerNormalization(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "SkipLayerNormalization", "SkipLayerNormalization", "add bias")
+
+    def fuse(self, node, input_name_to_nodes, output_name_to_node):
+        if len(node.input) != 4:
+            return
+
+        return_indice = []
+        nodes = self.model.match_parent_path(node, ["Add", "MatMul"], [None, None], output_name_to_node, return_indice)
+        if nodes is not None:
+            (add, _matmul) = nodes
+        else:
+            # In case of fp16, we could have a Cast between the MatMul and the bias Add
+            return_indice = []
+            nodes = self.model.match_parent_path(
+                node, ["Add", "Cast", "MatMul"], [None, None, None], output_name_to_node, return_indice
+            )
+            if nodes is not None:
+                (add, _cast, _matmul) = nodes
+            else:
+                return
+
+        assert len(return_indice) == 2 or len(return_indice) == 3
+        add_input_index = return_indice[0]
+        if add_input_index >= 2:
+            return
+        sln_input = add.input[return_indice[1]]
+        bias_input = add.input[1 - return_indice[1]]
+        skip_input = node.input[1 - add_input_index]
+
+        # bias should be one dimension
+        initializer = self.model.get_initializer(bias_input)
+        if initializer is None:
+            return
+        bias_weight = NumpyHelper.to_array(initializer)
+        if bias_weight is None:
+            logger.debug("Bias weight not found")
+            return
+        if len(bias_weight.shape) != 1:
+            logger.debug("Bias weight is not 1D")
+            return
+
+        subgraph_nodes = [node, add]
+        if not self.model.is_safe_to_fuse_nodes(subgraph_nodes, node.output, input_name_to_nodes, output_name_to_node):
+            logger.debug("Skip fusing SkipLayerNormalization with Bias since it is not safe")
+            return
+
+        self.nodes_to_remove.extend(subgraph_nodes)
+        inputs = [
+            sln_input,
+            skip_input,
+            node.input[2],
+            node.input[3],
+            bias_input,
+        ]
+        new_node = helper.make_node(
+            "SkipLayerNormalization",
+            inputs=inputs,
+            outputs=node.output,
+            name=self.model.create_node_name("SkipLayerNormalization", "SkipLayerNorm_AddBias_"),
+        )
+        new_node.domain = "com.microsoft"
+
+        # Pass attribute "epsilon" from skiplayernorm node to skiplayernorm(add bias)
+        for att in node.attribute:
+            if att.name == "epsilon":
+                new_node.attribute.extend([att])
+
+        # Set default epsilon if no epsilon exists from skiplayernorm
+        if len(new_node.attribute) == 0:
+            new_node.attribute.extend([helper.make_attribute("epsilon", 1.0e-12)])
+
+        self.nodes_to_add.append(new_node)
+        self.node_name_to_graph_name[new_node.name] = self.this_graph_name

onnxruntime/transformers/fusion_transpose.py

@@ -0,0 +1,167 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+
+from logging import getLogger
+
+from fusion_base import Fusion
+from fusion_utils import FusionUtils
+from onnx import NodeProto, TensorProto, helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionTranspose(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "Transpose", "Transpose")
+
+    def fuse(
+        self,
+        transpose_node: NodeProto,
+        input_name_to_nodes: dict[str, list[NodeProto]],
+        output_name_to_node: dict[str, NodeProto],
+    ):
+        """
+        Note that onnxruntime will do comprehensive transpose optimization after loading model.
+        The purpose of this fusion is to make graph clean before running onnxruntime.
+
+        Case 1:
+              (input)-->Transpose(perm=a)-->Transpose(perm=b)-->
+        After:
+              (input)-->Transpose(perm=a)-->  (this path can be removed if the output is not used anymore)
+                |
+                +----->Transpose(perm=a*b)-->
+
+        Case 2 (Cast has only one child):
+              (input)-->Transpose(perm=a)--> Cast -->Transpose(perm=b)-->
+        After:
+              (input)-->Transpose(perm=a)-->  (this path can be removed if the output is not used anymore)
+                |
+                +----->Cast --> Transpose(perm=a*b)-->
+        """
+        transpose_b = transpose_node
+        if transpose_b.input[0] not in output_name_to_node:
+            return
+
+        transpose_a = output_name_to_node[transpose_b.input[0]]
+        if transpose_a.op_type != "Cast":
+            cast_node = None
+        else:
+            cast_node = transpose_a
+
+            cast_children = self.model.get_children(cast_node, input_name_to_nodes)
+            if cast_children and len(cast_children) > 1:
+                return
+
+            if cast_node.input[0] not in output_name_to_node:
+                return
+
+            transpose_a = output_name_to_node[cast_node.input[0]]
+
+        if transpose_a.op_type != "Transpose":
+            return
+
+        permutation = OnnxModel.get_node_attribute(transpose_b, "perm")
+        assert isinstance(permutation, list)
+
+        parent_permutation = OnnxModel.get_node_attribute(transpose_a, "perm")
+        assert isinstance(parent_permutation, list)
+
+        assert len(parent_permutation) == len(permutation)
+
+        output_permutation = []
+        for _j, index in enumerate(permutation):
+            output_permutation.append(parent_permutation[index])
+
+        if cast_node is None:
+            if FusionUtils.skip_parent(self.model, transpose_b, transpose_a, input_name_to_nodes):
+                self.nodes_to_remove.append(transpose_a)
+        else:
+            if FusionUtils.skip_parent(self.model, cast_node, transpose_a, input_name_to_nodes):
+                self.nodes_to_remove.append(transpose_a)
+        transpose_b.ClearField("attribute")
+        transpose_b.attribute.extend([helper.make_attribute("perm", output_permutation)])
+
+
+class FusionInsertTranspose(Fusion):
+    def __init__(self, model: OnnxModel):
+        super().__init__(model, "", "GroupNorm")
+
+    def create_transpose_node(self, input_name: str, perm: list[int], output_name=None):
+        """Append a Transpose node after an input"""
+        node_name = self.model.create_node_name("Transpose")
+        if output_name is None:
+            output_name = node_name + "_out" + "-" + input_name
+        transpose_node = helper.make_node("Transpose", inputs=[input_name], outputs=[output_name], name=node_name)
+        transpose_node.attribute.extend([helper.make_attribute("perm", perm)])
+        return transpose_node
+
+    def fuse(
+        self,
+        group_norm_node: NodeProto,
+        input_name_to_nodes: dict[str, list[NodeProto]],
+        output_name_to_node: dict[str, NodeProto],
+    ):
+        """
+        This optimization will insert an Transpose, and onnxruntime transpose optimizer will remove it together with
+        another Transpose so that we can get effect of reducing one Transpose after onnxruntime optimization.
+        Before:
+            --> Gemm --> Unsqueeze(axes=[2]) --> Unsqueeze(axes=[3]) --> Add --> Transpose([0,2,3,1]) --> GroupNorm
+        After:
+            --> Gemm --> Unsqueeze(axes=[1]) --> Unsqueeze(axes=[2]) -->Transpose([0,3,1,2]) --> Add --> Transpose([0,2,3,1]) --> GroupNorm
+        """
+        gemm_path = self.model.match_parent_path(
+            group_norm_node, ["Transpose", "Add", "Unsqueeze", "Unsqueeze", "Gemm"], [0, 0, None, 0, 0]
+        )
+        if gemm_path is None:
+            return
+        transpose, add, unsqueeze_3, unsqueeze_2, gemm = gemm_path
+        if self.model.find_graph_output(unsqueeze_3.output[0]):
+            return
+
+        permutation = OnnxModel.get_node_attribute(transpose, "perm")
+        assert isinstance(permutation, list)
+        if permutation != [0, 2, 3, 1]:
+            return
+
+        if not (
+            len(unsqueeze_3.input) == 2
+            and self.model.get_constant_value(unsqueeze_3.input[1]) == 3
+            and len(unsqueeze_2.input) == 2
+            and self.model.get_constant_value(unsqueeze_2.input[1]) == 2
+            and len(self.model.get_children(gemm, input_name_to_nodes)) == 1
+            and len(self.model.get_children(unsqueeze_3, input_name_to_nodes)) == 1
+            and len(self.model.get_children(unsqueeze_2, input_name_to_nodes)) == 1
+        ):
+            return
+
+        # Here we use hard-coded name so that it could be shared for the whole model.
+        axes_1 = "ort_const_unsqueeze_axes_1"
+        if self.model.get_initializer(axes_1) is None:
+            self.add_initializer(
+                name=axes_1,
+                data_type=TensorProto.INT64,
+                dims=[1],
+                vals=[1],
+                raw=False,
+            )
+
+        axes_2 = "ort_const_unsqueeze_axes_2"
+        if self.model.get_initializer(axes_2) is None:
+            self.add_initializer(
+                name=axes_2,
+                data_type=TensorProto.INT64,
+                dims=[1],
+                vals=[2],
+                raw=False,
+            )
+
+        unsqueeze_3.input[1] = "ort_const_unsqueeze_axes_2"
+        unsqueeze_2.input[1] = "ort_const_unsqueeze_axes_1"
+        transpose_output_name = self.model.create_node_name("Transpose") + "_NCHW"
+        self.model.replace_input_of_all_nodes(unsqueeze_3.output[0], transpose_output_name)
+        new_transpose = self.create_transpose_node(unsqueeze_3.output[0], [0, 3, 1, 2], transpose_output_name)
+        self.model.add_node(new_transpose, self.this_graph_name)
+        self.increase_counter("Insert Transpose")

onnxruntime/transformers/fusion_utils.py

@@ -0,0 +1,321 @@
+# -------------------------------------------------------------------------
+# Copyright (c) Microsoft Corporation.  All rights reserved.
+# Licensed under the MIT License.
+# --------------------------------------------------------------------------
+from logging import getLogger
+
+import numpy
+from numpy import array_equal, ndarray
+from onnx import NodeProto, TensorProto, helper, numpy_helper
+from onnx_model import OnnxModel
+
+logger = getLogger(__name__)
+
+
+class FusionUtils:
+    def __init__(self, model: OnnxModel):
+        self.model: OnnxModel = model
+
+    def cast_graph_input_to_int32(self, input_name: str) -> tuple[bool, str]:
+        graph_input = self.model.find_graph_input(input_name)
+        if graph_input is not None and graph_input.type.tensor_type.elem_type != TensorProto.INT32:
+            cast_output, cast_node = self.cast_input_to_int32(input_name)
+            logger.debug(f"Casted graph input {input_name} to int32")
+            return True, cast_output
+
+        logger.debug(f"Did not cast graph input {input_name} to int32: found {graph_input is not None}")
+        return False, input_name
+
+    def cast_input(self, input_name: str, target_type="int32"):
+        output_name = input_name + "_" + target_type
+
+        if target_type == "int32":
+            to_type = int(TensorProto.INT32)
+        elif target_type == "float32":
+            to_type = int(TensorProto.FLOAT)
+        elif target_type == "float16":
+            to_type = int(TensorProto.FLOAT16)
+        else:
+            raise ValueError("Invalid target_type: {target_type}")
+
+        cast_node = self.add_cast_node(input_name, to_type, output_name)
+
+        return output_name, cast_node
+
+    def add_cast_node(
+        self,
+        input_name: str,
+        to_type: int,
+        output_name: str | None = None,
+        output_name_to_node=None,
+        graph_name: str | None = None,
+    ):
+        if output_name is None:
+            output_name = input_name + f"_cast_to_{to_type}"
+
+        # Avoid consequent Cast nodes.
+        inputs = [input_name]
+        if output_name_to_node is None:
+            output_name_to_node = self.model.output_name_to_node()
+        if input_name in output_name_to_node:
+            parent_node = output_name_to_node[input_name]
+            if parent_node and parent_node.op_type == "Cast":
+                inputs = [parent_node.input[0]]
+
+        cast_node = helper.make_node("Cast", inputs=inputs, outputs=[output_name])
+
+        cast_node.attribute.extend([helper.make_attribute("to", to_type)])
+        self.model.add_node(cast_node, graph_name=graph_name)
+
+        return cast_node
+
+    def cast_input_to_int32(self, input_name: str):
+        return self.cast_input(input_name, "int32")
+
+    def remove_cast_int32(self, input_name: str):
+        input_name_to_nodes = self.model.input_name_to_nodes()
+        nodes = input_name_to_nodes[input_name]
+        for node in nodes:
+            if node.op_type == "Cast":
+                is_int32 = False
+                for att in node.attribute:
+                    if att.name == "to" and att.i == int(TensorProto.INT32):
+                        is_int32 = True
+                        break
+                if is_int32:
+                    output_name = node.output[0]
+                    self.model.remove_node(node)
+                    self.model.replace_input_of_all_nodes(output_name, input_name)
+
+    @staticmethod
+    def update_node_input(node, i, new_input_name, input_name_to_nodes):
+        old_input_reference = 0
+        if (node.input[i] in input_name_to_nodes) and node in input_name_to_nodes[node.input[i]]:
+            input_name_to_nodes[node.input[i]].remove(node)
+            old_input_reference = len(input_name_to_nodes[node.input[i]])
+
+        node.input[i] = new_input_name
+
+        if new_input_name in input_name_to_nodes:
+            input_name_to_nodes[new_input_name].append(node)
+        else:
+            input_name_to_nodes[new_input_name] = [node]
+
+        return old_input_reference
+
+    @staticmethod
+    def skip_parent(model: OnnxModel, node, parent_node, input_name_to_nodes, node_input_index=0, parent_input_index=0):
+        """
+        Before:
+              (input)-->parent-->node-->(output)
+        After:
+              (input)-->parent-->
+                |
+                +----->node-->(output)
+
+        This function returns a flag whether the parent node can be removed.
+        """
+
+        old_input_name = node.input[node_input_index]
+        new_input_name = parent_node.input[parent_input_index]
+        old_input_reference = FusionUtils.update_node_input(node, node_input_index, new_input_name, input_name_to_nodes)
+
+        # We can remove the first Transpose if its output is not used (linked to graph output or other nodes) anymore.
+        parent_can_be_removed = (old_input_reference == 0) and not model.find_graph_output(old_input_name)
+
+        return parent_can_be_removed
+
+    def get_squeeze_or_unsqueeze_axes(self, node: NodeProto) -> ndarray | None:
+        assert node.op_type in ["Squeeze", "Unsqueeze"]
+
+        # For opset >= 13, axes is an input instead of an attribute.
+        if len(node.input) > 1:
+            return self.model.get_constant_value(node.input[1])
+
+        axes = None
+        for attr in node.attribute:
+            if attr.name == "axes":
+                axes = helper.get_attribute_value(attr)
+        return axes
+
+    @staticmethod
+    def check_node_attribute(node, attribute_name: str, expected_value, default_value=None):
+        """Verify that a node has expected value for an attribute.
+
+        Args:
+            node (NodeProto): a node to check
+            attribute_name (str): name of attribute
+            expected_value (Any): expected value of the attribute
+            default_value (Any, optional): default value if the attribute does not exist. Defaults to None.
+
+        Returns:
+            bool: whether the check is passed or not
+        """
+        value = default_value
+        for attr in node.attribute:
+            if attr.name == attribute_name:
+                value = helper.get_attribute_value(attr)
+
+        if isinstance(expected_value, list):
+            return (isinstance(value, (ndarray, list))) and array_equal(expected_value, value, equal_nan=False)
+        else:
+            return value == expected_value
+
+    @staticmethod
+    def transpose_2d_int8_tensor(tensor: TensorProto):
+        """Transpose a 2-D INT8 TensorProto
+        Args:
+            tensor (TensorProto): tensor to be transposed
+        Returns:
+            tensor (TensorProto): transposed tensor
+        """
+        if not isinstance(tensor, TensorProto):
+            raise TypeError(f"Expected input type is an ONNX TensorProto but got {type(tensor)}")
+
+        if len(tensor.dims) != 2 or tensor.data_type != TensorProto.INT8:
+            raise ValueError("Only INT8 2-D tensors can be transposed")
+
+        if tensor.raw_data:
+            int32_data = numpy.reshape(numpy.frombuffer(tensor.raw_data, dtype="int8"), tensor.dims)
+            int32_transposed_data = numpy.transpose(int32_data, [1, 0])
+            tensor.raw_data = int32_transposed_data.tobytes()
+
+        else:
+            raise ValueError("only raw buffer supported")
+
+        return tensor
+
+    @staticmethod
+    def check_qdq_node_for_fusion(node: NodeProto, model: OnnxModel, allow_per_tensor_quantization_only=True):
+        """Verify if a provided QuantizeLinear (Q) / DequantizeLinear (DQ) node is a good candidate for fusion.
+           It is a good candidate for fusion if:
+           (1) The Q/DQ node is for per-tensor quantization if allow_per_tensor_quantization_only is `True`
+           (2) The Q/DQ node should have constant scale
+           (3) The Q/DQ node should have a zero point of 0
+        Args:
+            node (NodeProto): a Q/DQ node to check
+        Returns:
+            bool: whether the check is passed or not
+        """
+        if node.op_type not in {"QuantizeLinear", "DequantizeLinear"}:
+            logger.debug(f"Provided node is not a Q/DQ node. Op Type: {node.op_type}")
+
+        scale = model.get_constant_value(node.input[1])
+
+        # Scale is not constant
+        if scale is None:
+            return False
+
+        # Not per-tensor quantization
+        scale_has_single_element = scale.ndim == 0 or (scale.ndim == 1 and scale.shape[0] == 1)
+        if allow_per_tensor_quantization_only and not scale_has_single_element:
+            return False
+
+        # If the Q/DQ node has no zero point input, it is assumed to be 0 (per ONNX spec)
+        if len(node.input) == 2:
+            return True
+
+        # Zero point should be constant and should have a value of 0
+        zero_point = model.get_constant_value(node.input[2])
+
+        # Zero point and scale should have same number of dims
+        if scale.ndim != zero_point.ndim:
+            return False
+
+        # Zero point is not constant or zero point is not zero
+        if zero_point is None:
+            return False
+
+        return numpy.all(zero_point == 0)
+
+    def check_node_input_value(self, node, input_index: int, expected_value):
+        """Verify that a node has expected input value
+
+        Args:
+            node (NodeProto): a node to check
+            input_index (int): index of its input to be verified
+            expected_value (Any): expected value of the input
+
+        Returns:
+            bool: whether the check is passed or not
+        """
+        assert len(node.input) > input_index
+
+        value = self.model.get_constant_value(node.input[input_index])
+
+        if isinstance(expected_value, list):
+            return (isinstance(value, (ndarray, list))) and array_equal(expected_value, value, equal_nan=False)
+        else:
+            return value == expected_value
+
+    def remove_identity_nodes(self):
+        """Remove Identity nodes, except those right before graph output."""
+        nodes_to_remove = []
+        graph_output_names = self.model.get_graphs_output_names()
+        for node in self.model.nodes():
+            if node.op_type == "Identity":
+                if node.output[0] not in graph_output_names:
+                    self.model.replace_input_of_all_nodes(node.output[0], node.input[0])
+                    nodes_to_remove.append(node)
+
+        if nodes_to_remove:
+            self.model.remove_nodes(nodes_to_remove)
+            logger.info(f"Removed {len(nodes_to_remove)} Identity nodes")
+
+    def remove_cascaded_cast_nodes(self):
+        self.model.remove_cascaded_cast_nodes()
+
+    def remove_useless_cast_nodes(self):
+        self.model.remove_useless_cast_nodes()
+
+    def remove_useless_reshape_nodes(self):
+        """Remove reshape node that is not needed based on symbolic shape inference: input and output has same shape"""
+        shape_infer = self.model.infer_runtime_shape(update=True)
+        if shape_infer is None:
+            return
+
+        nodes_to_remove = []
+        for node in self.model.nodes():
+            if node.op_type == "Reshape":
+                input_shape = shape_infer.get_edge_shape(node.input[0])
+                output_shape = shape_infer.get_edge_shape(node.output[0])
+                if input_shape and output_shape and input_shape == output_shape:
+                    logger.info(
+                        f"Remove reshape node {node.name} since its input shape is same as output: {input_shape}"
+                    )
+                    nodes_to_remove.append(node)
+
+        if nodes_to_remove:
+            graph_input_names = set(self.model.get_graphs_input_names())
+            graph_output_names = set(self.model.get_graphs_output_names())
+            for node in nodes_to_remove:
+                if bool(set(node.output) & graph_output_names):
+                    if (
+                        not bool(set(node.input) & graph_input_names)
+                        and len(self.model.input_name_to_nodes()[node.input[0]]) == 1  # parent has only one child
+                    ):
+                        self.model.replace_output_of_all_nodes(node.input[0], node.output[0])
+                    else:
+                        continue
+                else:
+                    self.model.replace_input_of_all_nodes(node.output[0], node.input[0])
+                self.model.remove_node(node)
+
+
+class NumpyHelper:
+    @staticmethod
+    def to_array(tensor: TensorProto, fill_zeros: bool = False) -> ndarray:
+        # When weights are in external data format but not presented, we can still test the optimizer with two changes:
+        # (1) set fill_zeros = True  (2) change load_external_data=False in optimizer.py
+        if fill_zeros:
+            return ndarray(
+                shape=tensor.dims,
+                dtype=helper.tensor_dtype_to_np_dtype(tensor.data_type),
+            )
+
+        if tensor.data_type == TensorProto.BFLOAT16:
+            import onnx_ir as ir  # noqa: PLC0415
+
+            # Use onnx_ir to correctly handle bfloat16 tensors
+            return ir.from_proto(tensor).numpy()
+        return numpy_helper.to_array(tensor)