onnxslim 0.1.68__py3-none-any.whl → 0.1.70__py3-none-any.whl

onnxslim/__init__.py

@@ -2,9 +2,7 @@ import os
 import warnings
 
 from onnxslim.cli import slim
-from onnxslim.core.optimization import OptimizationSettings
 from onnxslim.core.pattern.registry import (
-    DEFAULT_FUSION_PATTERNS,
     register_fusion_pattern,
 )
 from onnxslim.version import __version__

onnxslim/argparser.py

@@ -2,10 +2,28 @@ import argparse
 import dataclasses
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from dataclasses import dataclass, field
-from typing import List, Optional, Type, Union, get_args, get_origin
-
-import onnxslim
-
+from typing import List, Optional, Type, Union, get_args, get_origin, TypedDict, Dict, Literal
+
+from .core.optimization import OptimizationSettings
+from .core.pattern.registry import DEFAULT_FUSION_PATTERNS
+from .version import __version__
+
+
+class OnnxSlimKwargs(TypedDict, total=False):
+    model_check: bool
+    input_shapes: Dict[str, List[int]]
+    inputs: List[str]
+    outputs: List[str]
+    no_shape_infer: bool
+    skip_optimizations: List[str]
+    dtype: Literal["float16", "float32", "uint8", "int8"]
+    skip_fusion_patterns: List[str]
+    size_threshold: int
+    inspect: bool
+    dump_to_disk: bool
+    save_as_external_data: bool
+    model_check_inputs: Optional[List[str]]
+    verbose: bool
 
 def _get_inner_type(arg_type):
     if get_origin(arg_type) is Union:
@@ -42,14 +60,14 @@ class OptimizationArguments:
         default=None,
         metadata={
             "help": "whether to skip some optimizations",
-            "choices": list(onnxslim.OptimizationSettings.keys()),
+            "choices": list(OptimizationSettings.keys()),
         },
     )
     skip_fusion_patterns: Optional[List[str]] = field(
         default=None,
         metadata={
             "help": "whether to skip the fusion of some patterns",
-            "choices": list(onnxslim.DEFAULT_FUSION_PATTERNS.keys()),
+            "choices": list(DEFAULT_FUSION_PATTERNS.keys()),
         },
     )
     size_threshold: int = field(
@@ -173,7 +191,7 @@ class OnnxSlimArgumentParser(ArgumentParser):
         # Add positional arguments separately for ModelArguments
         self.parser.add_argument("input_model", help="input onnx model")
         self.parser.add_argument("output_model", nargs="?", default=None, help="output onnx model")
-        self.parser.add_argument("-v", "--version", action="version", version=onnxslim.__version__)
+        self.parser.add_argument("-v", "--version", action="version", version=__version__)
 
     def parse_args_into_dataclasses(self):
         # Pre-parse arguments to check for `--inspect`

onnxslim/cli/_main.py

@@ -2,8 +2,10 @@ from __future__ import annotations
 
 import onnx
 
+from onnxslim.argparser import OnnxSlimKwargs
 
-def slim(model: str | onnx.ModelProto | list[str | onnx.ModelProto], *args, **kwargs):
+
+def slim(model: str | onnx.ModelProto | list[str | onnx.ModelProto], *args, **kwargs: OnnxSlimKwargs):
     import os
     import time
     from pathlib import Path

onnxslim/third_party/symbolic_shape_infer.py

@@ -1944,7 +1944,25 @@ class SymbolicShapeInference:
 
     def _infer_Shape(self, node):  # noqa: N802
         """Infers and sets the symbolic shape for the output node in the computation graph."""
-        self.sympy_data_[node.output[0]] = self._get_sympy_shape(node, 0)
+        start = get_attribute(node, "start", 0)
+        end = get_attribute(node, "end", None)
+
+        full_sympy_shape = self._get_sympy_shape(node, 0)
+        num_dims = len(full_sympy_shape)
+
+        if start < 0:
+            start = num_dims + start
+        if end is None:
+            end = num_dims
+        elif end < 0:
+            end = num_dims + end
+
+        assert 0 <= start <= end <= num_dims, (
+            f"reshape start/end invalid: start={start}, end={end}, total_dims={num_dims}"
+        )
+
+        target_sympy_shape = full_sympy_shape[start:end]
+        self.sympy_data_[node.output[0]] = target_sympy_shape
 
     def _infer_Size(self, node):  # noqa: N802
         """Infers and sets the size of the output node by computing the product of its shape in the computation
@@ -2826,7 +2844,6 @@ class SymbolicShapeInference:
     def _infer_impl(self, start_sympy_data=None):
         """Infer implementation details and update symbolic data and input symbols."""
         self.sympy_data_ = start_sympy_data or {}
-        self.out_mp_.graph.ClearField("value_info")
         self._apply_suggested_merge(graph_input_only=True)
         self.input_symbols_ = set()
         for i in self.out_mp_.graph.input:

onnxslim/version.py

@@ -1 +1 @@
-__version__ = "0.1.68"
+__version__ = "0.1.70"

{onnxslim-0.1.68.dist-info → onnxslim-0.1.70.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: onnxslim
-Version: 0.1.68
+Version: 0.1.70
 Summary: OnnxSlim: A Toolkit to Help Optimize Onnx Model
 Home-page: https://github.com/inisis/OnnxSlim
 Author: inisis

{onnxslim-0.1.68.dist-info → onnxslim-0.1.70.dist-info}/RECORD

@@ -1,10 +1,10 @@
-onnxslim/__init__.py,sha256=-zmz0A4bd_t8VgNXZVMizIpKJswIlC99hr8juWYm2kc,660
+onnxslim/__init__.py,sha256=ECHGdxzg4b-4SZaPhxM_KulBi-xDbVcVUbpJc8i6a60,571
 onnxslim/__main__.py,sha256=FgDcl6xX8kV_52rB-jPVsmGqidlVhkpe_YhXK75-nFU,75
-onnxslim/argparser.py,sha256=aLVOV7c_HmsvuuAeb1FMRyqXTPRtjwm-bGt4O-7AxMg,8350
+onnxslim/argparser.py,sha256=pFv3nEZH2BiHO9ejS4Iq5ZuZ3GrpdyRQJypAyR0xF7w,8942
 onnxslim/utils.py,sha256=uNh3Mr4DberXtKL6p649QrrLR7gYbldYuIm6QLo14Zg,28043
-onnxslim/version.py,sha256=jLW9YDKJDUjNiMYmJhfPaIubYAn7SaPKsTXhnrOdbhg,23
+onnxslim/version.py,sha256=cccQoQHdRDqS_EuWR2yvG91WHNMC63UUtWzJGOrGCaY,23
 onnxslim/cli/__init__.py,sha256=kxK27cDgWotBOWRs86rbRQf_dtmniKr1GZJeasxfESE,42
-onnxslim/cli/_main.py,sha256=CPitq8khtOLCrJxvmam61ONppACfjQUzRshZtyabZak,5487
+onnxslim/cli/_main.py,sha256=jEKv9q7y_y9g0GsrfXcnk_wyMVej6jhe9QNPChE7yTs,5550
 onnxslim/core/__init__.py,sha256=d6zfD4rcR51yJIFg1LpGQt2tK4g1WLmA87_7Gna89OU,8469
 onnxslim/core/optimization/__init__.py,sha256=cAzv2YoZMS2krWoylrqDd9HU0aAhSCAM29dZclKqOQo,4777
 onnxslim/core/optimization/dead_node_elimination.py,sha256=xAO37JvW7Yb3mT5cXY6K6cWMAvbmFL0pczpcRQfd5Ek,7557
@@ -29,7 +29,7 @@ onnxslim/core/pattern/fusion/reduce.py,sha256=dMC7CPlFglrJxugsJWjcc-jQCIa_GIbW1y
 onnxslim/misc/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/misc/tabulate.py,sha256=pytOhVcWPW6B464BvG0rycbz8QA0MAeCSqQ7EO93tqs,99648
 onnxslim/third_party/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-onnxslim/third_party/symbolic_shape_infer.py,sha256=wG5wcB9TiUnOG_FX2w1XGy_oqKh2TMforFwtJrx_IV0,153088
+onnxslim/third_party/symbolic_shape_infer.py,sha256=iBrhPAgR--GNm6jSbR0gmmhVnzDICVULIdiqRlL7EIg,153592
 onnxslim/third_party/_sympy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/third_party/_sympy/functions.py,sha256=FrrTHPkDhwh_EccLGugbcodeOjfGeM17T5_oVPUjOlA,8877
 onnxslim/third_party/_sympy/numbers.py,sha256=w1dJJcQkKRzLDCJMn70YTwPtrEWPFrhCJpAsZhukJOk,11383
@@ -55,10 +55,10 @@ onnxslim/third_party/onnx_graphsurgeon/logger/logger.py,sha256=L12rrwn33RHH-2WLv
 onnxslim/third_party/onnx_graphsurgeon/util/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 onnxslim/third_party/onnx_graphsurgeon/util/exception.py,sha256=KrsHbKEQ4237UbjlODsUzvkXoAY72LZi23ApBeFANWg,786
 onnxslim/third_party/onnx_graphsurgeon/util/misc.py,sha256=kyxInD2SCRLU4wHMeiDEYEHB3871fGks6kQTuF9uATY,8960
-onnxslim-0.1.68.dist-info/licenses/LICENSE,sha256=oHZXw-yrBwdNVGu4JtlZhMgmQHKIZ7BJJlJdhu1HKvI,1062
-onnxslim-0.1.68.dist-info/METADATA,sha256=UiRpjeAqYUvRgvaNeSu6Na2IgmG1BNJmttF1s3uh_EU,7621
-onnxslim-0.1.68.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-onnxslim-0.1.68.dist-info/entry_points.txt,sha256=O2QgceCVeGeRhnxRSDRcGiFd0ZNfElwrTiRo1W2V7KA,47
-onnxslim-0.1.68.dist-info/top_level.txt,sha256=EWFTb99i0kc6cC9akqNKp88ipzg17_VZzYN7z1kQNlA,9
-onnxslim-0.1.68.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-onnxslim-0.1.68.dist-info/RECORD,,
+onnxslim-0.1.70.dist-info/licenses/LICENSE,sha256=oHZXw-yrBwdNVGu4JtlZhMgmQHKIZ7BJJlJdhu1HKvI,1062
+onnxslim-0.1.70.dist-info/METADATA,sha256=_Cv8cxA_t_EUzLq4rPTqnzhvvMs68Rs6exGDw1gxNVI,7621
+onnxslim-0.1.70.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+onnxslim-0.1.70.dist-info/entry_points.txt,sha256=O2QgceCVeGeRhnxRSDRcGiFd0ZNfElwrTiRo1W2V7KA,47
+onnxslim-0.1.70.dist-info/top_level.txt,sha256=EWFTb99i0kc6cC9akqNKp88ipzg17_VZzYN7z1kQNlA,9
+onnxslim-0.1.70.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+onnxslim-0.1.70.dist-info/RECORD,,