onnxslim 0.1.71__py3-none-any.whl → 0.1.72__py3-none-any.whl

onnxslim/core/__init__.py

@@ -1,6 +1,9 @@
+from __future__ import annotations
+
 import logging
 import os
 import tempfile
+from typing import Optional
 
 import numpy as np
 import onnx
@@ -18,6 +21,7 @@ logger = logging.getLogger("onnxslim")
 
 DEBUG = bool(os.getenv("ONNXSLIM_DEBUG"))
 AUTO_MERGE = True if os.getenv("ONNXSLIM_AUTO_MERGE") is None else bool(int(os.getenv("ONNXSLIM_AUTO_MERGE")))
+FORCE_ONNXRUNTIME_SHAPE_INFERENCE = bool(os.getenv("ONNXSLIM_FORCE_ONNXRUNTIME_SHAPE_INFERENCE"))
 
 
 def input_shape_modification(model: onnx.ModelProto, input_shapes: str) -> onnx.ModelProto:
@@ -122,6 +126,9 @@ def input_modification(model: onnx.ModelProto, inputs: str) -> onnx.ModelProto:
 def shape_infer(model: onnx.ModelProto):
     """Infer tensor shapes in an ONNX model using symbolic and static shape inference techniques."""
     logger.debug("Start shape inference.")
+    if FORCE_ONNXRUNTIME_SHAPE_INFERENCE:
+        logger.debug("force onnxruntime shape infer.")
+        return SymbolicShapeInference.infer_shapes(model, auto_merge=AUTO_MERGE)
     try:
         logger.debug("try onnxruntime shape infer.")
         model = SymbolicShapeInference.infer_shapes(model, auto_merge=AUTO_MERGE)
@@ -142,7 +149,7 @@ def shape_infer(model: onnx.ModelProto):
     return model
 
 
-def optimize(model: onnx.ModelProto, skip_fusion_patterns: str = None, size_threshold: int = None):
+def optimize(model: onnx.ModelProto, skip_fusion_patterns: str | None = None, size_threshold: int | None = None):
     """Optimize the given ONNX model with options to skip specific fusion patterns and return the optimized model."""
     logger.debug("Start converting model to gs.")
     graph = gs.import_onnx(model).toposort()
@@ -171,11 +178,11 @@ def convert_data_format(model: onnx.ModelProto, dtype: str) -> onnx.ModelProto:
 
         for node in graph.nodes:
             if node.op == "Cast":
-                inp_dtype = [input.dtype for input in node.inputs][0]
+                inp_dtype = next(input.dtype for input in node.inputs)
                 if inp_dtype in [np.float16, np.float32]:
                     node.erase()
                 else:
-                    outp_dtype = [output.dtype for output in node.outputs][0]
+                    outp_dtype = next(output.dtype for output in node.outputs)
                     if outp_dtype == np.float16:
                         node.attrs["to"] = dtype_to_onnx(np.float32)
                         node.outputs[0].dtype = np.float32

onnxslim/core/optimization/__init__.py

@@ -51,7 +51,7 @@ class OptimizationSettings:
         return any([getattr(cls, key) for key in cls.keys()])
 
 
-def optimize_model(model: onnx.ModelProto | gs.Graph, skip_fusion_patterns: str = None) -> onnx.ModelProto:
+def optimize_model(model: onnx.ModelProto | gs.Graph, skip_fusion_patterns: str | None = None) -> onnx.ModelProto:
     """Optimize and transform the given ONNX model using various fusion patterns and graph rewriting techniques."""
     graph = model if isinstance(model, gs.Graph) else gs.import_onnx(model)
     if OptimizationSettings.graph_fusion:
@@ -85,7 +85,7 @@ def replace_custom_layer(
     inputs,
     outputs: list[str],
     name: str,
-    attrs: dict = None,
+    attrs: dict | None = None,
     domain: str = "ai.onnx.contrib",
 ):
     """Replace a custom layer in the computational graph with specified parameters and domain."""

onnxslim/core/optimization/dead_node_elimination.py

@@ -29,7 +29,7 @@ def dead_node_elimination(graph, is_subgraph=False):
                     node.erase()
                     logger.debug(f"removing {node.op} op: {node.name}")
         elif node.op == "Cast":
-            inp_dtype = [dtype_to_onnx(input.dtype) for input in node.inputs][0]
+            inp_dtype = next(dtype_to_onnx(input.dtype) for input in node.inputs)
             if inp_dtype == node.attrs["to"]:
                 node.erase()
                 logger.debug(f"removing {node.op} op: {node.name}")

onnxslim/core/pattern/__init__.py

@@ -114,10 +114,17 @@ class PatternMatcher:
         """Retrieve the match point node from the pattern dictionary based on output node input names."""
         return self.pattern_dict[self.pattern_dict[self.output_names[0]].input_names[0]]
 
+    def reset_input(self):
+        for k, v in self.pattern_dict.items():
+            if v.op == "input":
+                if hasattr(self, v.name):
+                    delattr(self, v.name)
+
     def match(self, node):
         """Match a given node to a pattern by comparing input names with the match point node from the pattern
         dictionary.
         """
+        self.reset_input()
         match_point = self.get_match_point()
 
         def match_(node, pattern_node):
@@ -125,8 +132,14 @@ class PatternMatcher:
             dictionary.
             """
             if pattern_node.op == "input":
-                return True
-
+                if hasattr(self, pattern_node.name):
+                    if getattr(self, pattern_node.name) == node:
+                        return True
+                    else:
+                        return False
+                else:
+                    setattr(self, pattern_node.name, node)
+                    return True
             # node is an input variable
             if not hasattr(node, "op"):
                 return False

onnxslim/core/pattern/elimination/slice.py

@@ -58,7 +58,7 @@ class SlicePatternMatcher(PatternMatcher):
                     inputs = []
                     inputs.extend(
                         (
-                            list(first_slice_node.inputs)[0],
+                            next(iter(first_slice_node.inputs)),
                             gs.Constant(
                                 second_slice_node_inputs[1].name + "_starts",
                                 values=np.array(new_starts, dtype=np.int64),

onnxslim/core/pattern/elimination/unsqueeze.py

@@ -50,14 +50,14 @@ class UnsqueezePatternMatcher(PatternMatcher):
                     axis + sum(1 for axis_ in axes_node_unsqueeze_1 if axis_ <= axis) for axis in axes_node_unsqueeze_0
                 ]
 
-                inputs = [list(node_unsqueeze_0.inputs)[0]]
+                inputs = [next(iter(node_unsqueeze_0.inputs))]
                 outputs = list(node_unsqueeze_1.outputs)
 
                 index = node_unsqueeze_1.inputs.index(node_unsqueeze_0.outputs[0])
                 node_unsqueeze_1.inputs.pop(index)
                 for i, item in enumerate(node_unsqueeze_0.inputs):
                     node_unsqueeze_1.inputs.insert(index + i, item)
-                inputs = [list(node_unsqueeze_1.inputs)[0]]
+                inputs = [next(iter(node_unsqueeze_1.inputs))]
                 outputs = list(node_unsqueeze_1.outputs)
                 node_unsqueeze_1.inputs.clear()
                 node_unsqueeze_1.outputs.clear()

onnxslim/core/pattern/fusion/concat_reshape.py

@@ -36,7 +36,7 @@ class ConcatReshapeMatcher(PatternMatcher):
     def rewrite(self, opset=11):
         match_case = {}
         concat_node = self.concat_0
-        index = [idx for idx, i in enumerate(concat_node.inputs) if isinstance(i, gs.Variable)][0]
+        index = next(idx for idx, i in enumerate(concat_node.inputs) if isinstance(i, gs.Variable))
         constant = gs.Constant(
             concat_node.inputs[index].name + "_fixed",
             values=np.array([-1], dtype=np.int64),

onnxslim/core/pattern/fusion/convadd.py

@@ -41,7 +41,7 @@ class ConvAddMatcher(PatternMatcher):
                 conv_bias = conv_node.inputs[2].values + node.inputs[1].values.squeeze()
 
             inputs = []
-            inputs.append(list(conv_node.inputs)[0])
+            inputs.append(next(iter(conv_node.inputs)))
             inputs.append(conv_weight)
             weight_name = list(conv_node.inputs)[1].name
             if weight_name.endswith("weight"):

onnxslim/core/pattern/fusion/convbn.py

@@ -53,7 +53,7 @@ class ConvBatchNormMatcher(PatternMatcher):
             conv_b = (conv_transpose_bias - bn_running_mean) * bn_var_rsqrt * bn_scale + bn_bias
 
             inputs = []
-            inputs.append(list(conv_transpose_node.inputs)[0])
+            inputs.append(next(iter(conv_transpose_node.inputs)))
             weight_name = list(conv_transpose_node.inputs)[1].name
             if weight_name.endswith("weight"):
                 bias_name = f"{weight_name[:-6]}bias"

onnxslim/core/pattern/fusion/gemm.py

@@ -105,7 +105,7 @@ class MatMulAddPatternMatcher(PatternMatcher):
                     }
                 )
 
-                values = list(input_variable.shape[:-1]) + [matmul_bias_variable.values.shape[-1]]
+                values = [*list(input_variable.shape[:-1]), matmul_bias_variable.values.shape[-1]]
                 post_reshape_const = gs.Constant(
                     f"{matmul_node.name}_post_reshape_in",
                     values=np.array(values, dtype=np.int64),

onnxslim/core/pattern/fusion/padconv.py

@@ -37,8 +37,7 @@ class PadConvMatcher(PatternMatcher):
 
         pad_inputs = len(pad_node.inputs)
         if pad_inputs < 3 or (
-            pad_inputs >= 3
-            and (isinstance(pad_node.inputs[2], gs.Constant) and pad_node.inputs[2].values == 0)
+            (pad_inputs >= 3 and (isinstance(pad_node.inputs[2], gs.Constant) and pad_node.inputs[2].values == 0))
             or (pad_inputs >= 3 and (isinstance(pad_node.inputs[2], gs.Variable) and pad_node.inputs[2].name == ""))
         ):
             if (

onnxslim/core/pattern/registry.py

@@ -1,3 +1,5 @@
+from __future__ import annotations
+
 from collections import OrderedDict
 
 DEFAULT_FUSION_PATTERNS = OrderedDict()
@@ -12,7 +14,7 @@ def register_fusion_pattern(fusion_pattern):
     DEFAULT_FUSION_PATTERNS[layer_type] = fusion_pattern
 
 
-def get_fusion_patterns(skip_fusion_patterns: str = None):
+def get_fusion_patterns(skip_fusion_patterns: str | None = None):
     """Returns a copy of the default fusion patterns, optionally excluding specific patterns."""
     default_fusion_patterns = DEFAULT_FUSION_PATTERNS.copy()
     if skip_fusion_patterns:

onnxslim/misc/tabulate.py

@@ -24,7 +24,7 @@ def _is_file(f):
     return isinstance(f, io.IOBase)
 
 
-__all__ = ["tabulate", "tabulate_formats", "simple_separated_format"]
+__all__ = ["simple_separated_format", "tabulate", "tabulate_formats"]
 try:
     from .version import version as __version__  # noqa: F401
 except ImportError:
@@ -1202,7 +1202,7 @@ def _format(val, valtype, floatfmt, intfmt, missingval="", has_invisible=True):
     if val is None:
         return missingval
 
-    if valtype is str or valtype is not int and valtype is not bytes and valtype is not float:
+    if valtype is str or (valtype is not int and valtype is not bytes and valtype is not float):
         return f"{val}"
     elif valtype is int:
         return format(val, intfmt)
@@ -1276,7 +1276,7 @@ def _prepend_row_index(rows, index):
     index_iter = iter(index)
     for row in sans_rows:
         index_v = next(index_iter)
-        new_rows.append([index_v] + list(row))
+        new_rows.append([index_v, *list(row)])
     rows = new_rows
     _reinsert_separating_lines(rows, separating_lines)
     return rows
@@ -1325,7 +1325,7 @@ def _normalize_tabular_data(tabular_data, headers, showindex="default"):
     """
     try:
         bool(headers)
-        is_headers2bool_broken = False  # noqa
+        is_headers2bool_broken = False
     except ValueError:  # numpy.ndarray, pandas.core.index.Index, ...
         is_headers2bool_broken = True  # noqa
         headers = list(headers)
@@ -1423,7 +1423,7 @@ def _normalize_tabular_data(tabular_data, headers, showindex="default"):
     if headers == "firstrow":
         if len(rows) > 0:
             if index is not None:
-                headers = [index[0]] + list(rows[0])
+                headers = [index[0], *list(rows[0])]
                 index = index[1:]
             else:
                 headers = rows[0]
@@ -2534,8 +2534,10 @@ class _CustomTextWrap(textwrap.TextWrapper):
                 if (
                     self.max_lines is None
                     or len(lines) + 1 < self.max_lines
-                    or (not chunks or self.drop_whitespace and len(chunks) == 1 and not chunks[0].strip())
-                    and cur_len <= width
+                    or (
+                        (not chunks or (self.drop_whitespace and len(chunks) == 1 and not chunks[0].strip()))
+                        and cur_len <= width
+                    )
                 ):
                     # Convert current line back to a string and store it in
                     # list of all lines (return value).

onnxslim/third_party/_sympy/printers.py

@@ -1,5 +1,6 @@
+from __future__ import annotations
+
 import sys
-from typing import Optional
 
 import sympy
 from sympy.printing.precedence import PRECEDENCE, precedence
@@ -22,7 +23,7 @@ class ExprPrinter(StrPrinter):
     def _print_Not(self, expr: sympy.Expr) -> str:
         return f"not ({self._print(expr.args[0])})"
 
-    def _print_Add(self, expr: sympy.Expr, order: Optional[str] = None) -> str:
+    def _print_Add(self, expr: sympy.Expr, order: str | None = None) -> str:
         return self.stringify(expr.args, " + ", precedence(expr))
 
     def _print_Relational(self, expr: sympy.Expr) -> str:

onnxslim/third_party/_sympy/symbol.py

@@ -12,9 +12,10 @@ You can occasionally test if prefixes have been hardcoded by renaming prefixes
 in this file and seeing what breaks.
 """
 
+from __future__ import annotations
+
 from collections.abc import Iterable
 from enum import Enum, auto
-from typing import Union
 
 import sympy
 
@@ -88,7 +89,7 @@ def make_symbol(prefix: SymT, idx: int, **kwargs) -> sympy.Symbol:
 
 # This type is a little wider than it should be, because free_symbols says
 # that it contains Basic, rather than Symbol
-def symbol_is_type(sym: sympy.Basic, prefix: Union[SymT, Iterable[SymT]]) -> bool:
+def symbol_is_type(sym: sympy.Basic, prefix: SymT | Iterable[SymT]) -> bool:
     assert isinstance(sym, sympy.Symbol)
     name_str = sym.name.lower()  # Match capitalized names like XBLOCK, RBLOCK
     if isinstance(prefix, SymT):
@@ -97,5 +98,5 @@ def symbol_is_type(sym: sympy.Basic, prefix: Union[SymT, Iterable[SymT]]) -> boo
         return name_str.startswith(tuple(prefix_str[p] for p in prefix))
 
 
-def free_symbol_is_type(e: sympy.Expr, prefix: Union[SymT, Iterable[SymT]]) -> bool:
+def free_symbol_is_type(e: sympy.Expr, prefix: SymT | Iterable[SymT]) -> bool:
     return any(symbol_is_type(v, prefix) for v in e.free_symbols)

onnxslim/third_party/onnx_graphsurgeon/exporters/onnx_exporter.py

@@ -278,8 +278,8 @@ class OnnxExporter(BaseExporter):
     def export_graph(
         graph_proto: onnx.GraphProto,
         graph: Graph,
-        tensor_map: OrderedDict[str, Tensor] = None,
-        subgraph_tensor_map: OrderedDict[str, Tensor] = None,
+        tensor_map: OrderedDict[str, Tensor] | None = None,
+        subgraph_tensor_map: OrderedDict[str, Tensor] | None = None,
         do_type_check=True,
     ) -> None:
         """

onnxslim/third_party/onnx_graphsurgeon/graph_pattern/graph_pattern.py

@@ -82,7 +82,7 @@ class PatternMapping(dict):
     def __str__(self) -> str:
         """Returns a string representation of the pattern mapping, including inputs, outputs, and constants."""
         if self.onnx_node is None:
-            return "{" + str.join(", ", [f"{key}: {str(value)}" for key, value in self.items()]) + "}"
+            return "{" + str.join(", ", [f"{key}: {value!s}" for key, value in self.items()]) + "}"
         return self.onnx_node.name
 
 

onnxslim/third_party/onnx_graphsurgeon/importers/onnx_importer.py

@@ -357,7 +357,7 @@ class OnnxImporter(BaseImporter):
     @staticmethod
     def import_function(
         onnx_function: onnx.FunctionProto,
-        model_opset: int = None,
+        model_opset: int | None = None,
         model_import_domains: onnx.OperatorSetIdProto = None,
     ) -> Function:
         """Imports an ONNX function to a Function object using the model opset and import domains."""
@@ -405,13 +405,13 @@ class OnnxImporter(BaseImporter):
     @staticmethod
     def import_graph(
         onnx_graph: onnx.GraphProto,
-        tensor_map: OrderedDict[str, Tensor] = None,
+        tensor_map: OrderedDict[str, Tensor] | None = None,
         opset=None,
         import_domains: onnx.OperatorSetIdProto = None,
         ir_version=None,
-        producer_name: str = None,
-        producer_version: str = None,
-        functions: list[Function] = None,
+        producer_name: str | None = None,
+        producer_version: str | None = None,
+        functions: list[Function] | None = None,
         metadata_props=None,
     ) -> Graph:
         """

onnxslim/third_party/onnx_graphsurgeon/ir/function.py

@@ -43,15 +43,15 @@ class Function(Graph):
     def __init__(
         self,
         name: str,
-        domain: str = None,
-        nodes: Sequence[Node] = None,
-        inputs: Sequence[Tensor] = None,
-        outputs: Sequence[Tensor] = None,
-        doc_string: str = None,
-        opset: int = None,
-        import_domains: Sequence[onnx.OperatorSetIdProto] = None,
-        functions: Sequence[Function] = None,
-        attrs: dict = None,
+        domain: str | None = None,
+        nodes: Sequence[Node] | None = None,
+        inputs: Sequence[Tensor] | None = None,
+        outputs: Sequence[Tensor] | None = None,
+        doc_string: str | None = None,
+        opset: int | None = None,
+        import_domains: Sequence[onnx.OperatorSetIdProto] | None = None,
+        functions: Sequence[Function] | None = None,
+        attrs: dict | None = None,
     ):
         """
         Args:

onnxslim/third_party/onnx_graphsurgeon/ir/graph.py

@@ -100,17 +100,17 @@ class Graph:
 
     def __init__(
         self,
-        nodes: Sequence[Node] = None,
-        inputs: Sequence[Tensor] = None,
-        outputs: Sequence[Tensor] = None,
+        nodes: Sequence[Node] | None = None,
+        inputs: Sequence[Tensor] | None = None,
+        outputs: Sequence[Tensor] | None = None,
         name=None,
         doc_string=None,
         opset=None,
         import_domains=None,
         ir_version=None,
-        producer_name: str = None,
-        producer_version: str = None,
-        functions: Sequence[Function] = None,
+        producer_name: str | None = None,
+        producer_version: str | None = None,
+        functions: Sequence[Function] | None = None,
         metadata_props=None,
     ):
         """
@@ -347,7 +347,7 @@ class Graph:
                     func_ids.add(func.unique_id)
             return self.functions
 
-        for graph in self.functions + [self]:
+        for graph in [*self.functions, self]:
             for subgraph in graph.subgraphs(recursive=True):
                 new_list = absorb_function_list(subgraph.functions)
                 subgraph._functions = new_list
@@ -774,7 +774,7 @@ class Graph:
                     if len(node.attrs) != 1:
                         G_LOGGER.warning("Constant node must contain exactly one attribute")
                         continue
-                    attr_name, attr_val = list(node.attrs.items())[0]
+                    attr_name, attr_val = next(iter(node.attrs.items()))
                     allowed_attrs = {
                         "value",
                         "value_float",
@@ -975,7 +975,7 @@ class Graph:
 
         def get_scalar_value(tensor):
             """Gets the scalar value of a constant tensor with a single item."""
-            return list(tensor.values)[0] if tensor.shape else tensor.values
+            return next(iter(tensor.values)) if tensor.shape else tensor.values
 
         def fold_shape(tensor):
             """Returns the input tensor shape if available, otherwise returns None."""
@@ -1409,7 +1409,7 @@ class Graph:
         self.nodes.append(node)
         return node.outputs
 
-    def copy(self, tensor_map: OrderedDict[str, Tensor] = None):
+    def copy(self, tensor_map: OrderedDict[str, Tensor] | None = None):
         """
         Copy the graph.
 

onnxslim/third_party/onnx_graphsurgeon/ir/node.py

@@ -42,11 +42,11 @@ class Node:
     def __init__(
         self,
         op: str,
-        name: str = None,
-        attrs: dict[str, object] = None,
-        inputs: list[Tensor] = None,
-        outputs: list[Tensor] = None,
-        domain: str = None,
+        name: str | None = None,
+        attrs: dict[str, object] | None = None,
+        inputs: list[Tensor] | None = None,
+        outputs: list[Tensor] | None = None,
+        domain: str | None = None,
     ):
         """
         A node represents an operation in a graph, and consumes zero or more Tensors, and produces zero or more Tensors.
@@ -152,8 +152,8 @@ class Node:
 
     def copy(
         self,
-        inputs: list[Tensor] = None,
-        outputs: list[Tensor] = None,
+        inputs: list[Tensor] | None = None,
+        outputs: list[Tensor] | None = None,
         tensor_map=None,
     ):
         """

onnxslim/third_party/onnx_graphsurgeon/ir/tensor.py

@@ -65,7 +65,7 @@ class Tensor:
     def to_constant(
         self,
         values: np.ndarray,
-        data_location: int = None,
+        data_location: int | None = None,
         export_dtype: np.dtype | onnx.TensorProto.DataType = None,
     ):
         """
@@ -91,7 +91,7 @@ class Tensor:
 
         return self
 
-    def to_variable(self, dtype: np.dtype | onnx.TensorProto.DataType = None, shape: Sequence[int | str] = None):
+    def to_variable(self, dtype: np.dtype | onnx.TensorProto.DataType = None, shape: Sequence[int | str] | None = None):
         """
         Modifies this tensor in-place to convert it to a Variable. This means that all consumers/producers of the tensor
         will see the update.
@@ -199,7 +199,7 @@ class Variable(Tensor):
         self,
         name: str,
         dtype: np.dtype | onnx.TensorProto.DataType = None,
-        shape: Sequence[int | str] = None,
+        shape: Sequence[int | str] | None = None,
         type: str = "tensor_type",
     ):
         """
@@ -390,7 +390,7 @@ class Constant(Tensor):
         self,
         name: str,
         values: np.ndarray | LazyValues,
-        data_location: int = None,
+        data_location: int | None = None,
         export_dtype: np.dtype | onnx.TensorProto.DataType = None,
     ):
         """
@@ -426,7 +426,7 @@ class Constant(Tensor):
         self.data_location = data_location
         self._export_dtype = export_dtype
 
-    def to_variable(self, dtype: np.dtype = None, shape: Sequence[int | str] = None):
+    def to_variable(self, dtype: np.dtype = None, shape: Sequence[int | str] | None = None):
         """Convert instance values to an appropriate variable with specified dtype and shape."""
         if shape is None:
             shape = []