onnxslim 0.1.81__py3-none-any.whl → 0.1.83__py3-none-any.whl

onnxslim/core/optimization/dead_node_elimination.py

@@ -54,9 +54,16 @@ def dead_node_elimination(graph, is_subgraph=False):
                     node.inputs.insert(1, reshape_const)
                     logger.debug(f"replacing {node.op} op: {node.name}")
         elif node.op == "Slice":
-            if node.inputs[0].shape and node.outputs[0].shape and node.inputs[0].shape == node.outputs[0].shape:
-                node.erase()
-                logger.debug(f"removing {node.op} op: {node.name}")
+            if (node.inputs[0].shape and node.outputs[0].shape
+                and node.inputs[0].shape == node.outputs[0].shape
+                and all(isinstance(item, int) for item in node.inputs[0].shape)):
+
+                # Check if slice is a no-op by analyzing parameters directly
+                # Slice inputs: data, starts, ends, [axes], [steps]
+                if is_noop_slice(node):
+                    node.erase()
+                    logger.debug(f"removing {node.op} op: {node.name}")
+
         elif node.op == "Mul":
             if (isinstance(node.inputs[1], Constant) and isinstance(node.inputs[0], Variable)) or (
                 isinstance(node.inputs[0], Constant) and isinstance(node.inputs[1], Variable)
@@ -153,3 +160,77 @@ def get_constant_variable(node, return_idx=False):
     for idx, input in enumerate(list(node.inputs)):
         if isinstance(input, Constant):
             return (idx, input) if return_idx else input
+
+
+def is_noop_slice(node):
+    """Check if a Slice node is a no-op by analyzing its parameters directly.
+
+    A Slice is a no-op when it extracts the entire tensor, i.e., for each sliced axis:
+    - start == 0 (or equivalent negative index)
+    - end >= dim_size (or is INT_MAX-like value)
+    - step == 1
+    """
+    # Slice inputs: data, starts, ends, [axes], [steps]
+    if len(node.inputs) < 3:
+        return False
+
+    data_shape = node.inputs[0].shape
+    if not data_shape or not all(isinstance(d, int) for d in data_shape):
+        return False
+
+    # Get starts and ends (required)
+    starts_input = node.inputs[1]
+    ends_input = node.inputs[2]
+
+    if not isinstance(starts_input, Constant) or not isinstance(ends_input, Constant):
+        return False
+
+    starts = starts_input.values.flatten().tolist()
+    ends = ends_input.values.flatten().tolist()
+
+    # Get axes (optional, defaults to [0, 1, 2, ...])
+    if len(node.inputs) > 3 and isinstance(node.inputs[3], Constant):
+        axes = node.inputs[3].values.flatten().tolist()
+    else:
+        axes = list(range(len(starts)))
+
+    # Get steps (optional, defaults to [1, 1, 1, ...])
+    if len(node.inputs) > 4 and isinstance(node.inputs[4], Constant):
+        steps = node.inputs[4].values.flatten().tolist()
+    else:
+        steps = [1] * len(starts)
+
+    # Check each axis
+    ndim = len(data_shape)
+    for start, end, axis, step in zip(starts, ends, axes, steps):
+        # Normalize negative axis
+        if axis < 0:
+            axis = ndim + axis
+
+        if axis < 0 or axis >= ndim:
+            return False
+
+        dim_size = data_shape[axis]
+
+        # Step must be 1 for no-op
+        if step != 1:
+            return False
+
+        # Normalize negative start index
+        if start < 0:
+            start = max(0, dim_size + start)
+
+        # Start must be 0
+        if start != 0:
+            return False
+
+        # Normalize negative end index
+        if end < 0:
+            end = dim_size + end
+
+        # End must cover the entire dimension
+        # Common patterns: end == dim_size, or end is a large value like INT_MAX
+        if end < dim_size:
+            return False
+
+    return True

onnxslim/core/pattern/fusion/convadd.py

@@ -27,12 +27,13 @@ class ConvAddMatcher(PatternMatcher):
         conv_weight = list(conv_node.inputs)[1]
         conv_node_users = conv_node.users
         node = self.add_0
+        oc_axis = 0 if conv_node.op == "Conv" else 1 # output_channel_axis
         if (
             len(conv_node_users) == 1
             and isinstance(node.inputs[1], gs.Constant)
             and isinstance(conv_weight, gs.Constant)
             and node.inputs[1].values.squeeze().ndim == 1
-            and node.inputs[1].values.squeeze().shape[0] == conv_weight.shape[0]
+            and node.inputs[1].values.squeeze().shape[0] == conv_weight.shape[oc_axis]
         ):
             add_node = node
             if len(conv_node.inputs) == 2:
@@ -66,5 +67,24 @@ class ConvAddMatcher(PatternMatcher):
 
         return match_case
 
+class ConvTransposeAddMatcher(ConvAddMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvTransposeAddMatcher for fusing ConvTranspose and Add layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input             input  0  1 conv_0
+            ConvTranspose     conv_0 1+ 1 input bn_0
+            Add               add_0  2  1 conv_0 ? output
+            output            output 1  0 add_0
+            """
+        )
+        super(ConvAddMatcher, self).__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvTransposeAdd pattern."""
+        return "FusionConvTransposeAdd"
+
 
 register_fusion_pattern(ConvAddMatcher(1))
+register_fusion_pattern(ConvTransposeAddMatcher(1))

onnxslim/core/pattern/fusion/convbn.py

@@ -44,11 +44,9 @@ class ConvBatchNormMatcher(PatternMatcher):
                 conv_transpose_bias = conv_transpose_node.inputs[2].values
 
             bn_var_rsqrt = bn_scale / np.sqrt(bn_running_var + bn_eps)
+            oc_axis = 0 if conv_transpose_node.op == "Conv" else 1 # output_channel_axis
             shape = [1] * len(conv_transpose_weight.shape)
-            if bn_node.i(0).op == "Conv":
-                shape[0] = -1
-            else:
-                shape[1] = -1
+            shape[oc_axis] = -1
             conv_w = conv_transpose_weight * bn_var_rsqrt.reshape(shape)
             conv_b = (conv_transpose_bias - bn_running_mean) * bn_var_rsqrt + bn_bias
 
@@ -82,5 +80,24 @@ class ConvBatchNormMatcher(PatternMatcher):
 
         return match_case
 
+class ConvTransposeBatchNormMatcher(ConvBatchNormMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvTransposeBatchNormMatcher for fusing ConvTranspose and BatchNormalization layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input              input  0  1 conv_0
+            ConvTranspose      conv_0 1+ 1 input bn_0
+            BatchNormalization bn_0   5  1 conv_0 ? ? ? ? output
+            output             output 1  0 bn_0
+            """
+        )
+        super(ConvBatchNormMatcher, self).__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvTransposeBN pattern."""
+        return "FusionConvTransposeBN"
+
 
 register_fusion_pattern(ConvBatchNormMatcher(1))
+register_fusion_pattern(ConvTransposeBatchNormMatcher(1))

onnxslim/core/pattern/fusion/convmul.py

@@ -28,11 +28,10 @@ class ConvMulMatcher(PatternMatcher):
         conv_weight = list(conv_node.inputs)[1]
         if len(conv_node.users) == 1 and conv_node.users[0] == mul_node and isinstance(mul_node.inputs[1], gs.Constant):
             mul_constant = mul_node.inputs[1].values
-
-            if mul_constant.squeeze().ndim == 1 and mul_constant.squeeze().shape[0] == conv_weight.shape[0]:
-                weight_shape = conv_weight.values.shape
-                reshape_shape = [-1] + [1] * (len(weight_shape) - 1)
-
+            oc_axis = 0 if conv_node.op == "Conv" else 1 # output_channel_axis
+            if mul_constant.squeeze().ndim == 1 and mul_constant.squeeze().shape[0] == conv_weight.shape[oc_axis]:
+                reshape_shape = [1] * len(conv_weight.values.shape)
+                reshape_shape[oc_axis] = -1
                 mul_scale_reshaped = mul_constant.squeeze().reshape(reshape_shape)
                 new_weight = conv_weight.values * mul_scale_reshaped
 
@@ -65,5 +64,24 @@ class ConvMulMatcher(PatternMatcher):
 
         return match_case
 
+class ConvTransposeMulMatcher(ConvMulMatcher):
+    def __init__(self, priority):
+        """Initializes the ConvTransposeMulMatcher for fusing ConvTranspose and Mul layers in an ONNX graph."""
+        pattern = Pattern(
+            """
+            input            input  0  1 conv_0
+            ConvTranspose    conv_0 1+ 1 input mul_0
+            Mul              mul_0  2  1 conv_0 ? output
+            output           output 1  0 mul_0
+            """
+        )
+        super(ConvMulMatcher, self).__init__(pattern, priority)
+
+    @property
+    def name(self):
+        """Returns the name of the FusionConvTransposeMul pattern."""
+        return "FusionConvTransposeMul"
+
 
 register_fusion_pattern(ConvMulMatcher(1))
+register_fusion_pattern(ConvTransposeMulMatcher(1))

onnxslim/core/pattern/fusion/padconv.py

@@ -37,6 +37,8 @@ class PadConvMatcher(PatternMatcher):
         pad_node_users = pad_node.users
 
         pad_inputs = len(pad_node.inputs)
+        auto_pad = pad_node.attrs.get("auto_pad", "NOTSET")
+
         if pad_inputs < 3 or (
             (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 == ""))
@@ -45,6 +47,7 @@ class PadConvMatcher(PatternMatcher):
                 isinstance(pad_node.inputs[1], gs.Constant)
                 and pad_node.attrs.get("mode", "constant") == "constant"
                 and conv_node.inputs[1].shape
+                and (auto_pad == "NOTSET" or auto_pad == "VALID")
             ):
                 conv_weight_dim = len(conv_node.inputs[1].shape)
                 pad_value = pad_node.inputs[1].values.tolist()
@@ -74,6 +77,8 @@ class PadConvMatcher(PatternMatcher):
                         pads = [pad + conv_pad for pad, conv_pad in zip(pads, conv_pads)]
 
                     attrs["pads"] = pads
+                    conv_node.attrs.pop("auto_pad", None)
+
                     match_case[conv_node.name] = {
                         "op": "Conv",
                         "inputs": inputs,

onnxslim/core/shape_inference/__init__.py

@@ -0,0 +1,378 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""
+Symbolic Shape Inference Module
+
+This module provides symbolic shape inference for ONNX models. It replaces the
+monolithic SymbolicShapeInference class with a modular, handler-based architecture.
+
+Usage:
+    from onnxslim.core.shape_inference import ShapeInferencer
+
+    model = onnx.load("model.onnx")
+    model_with_shapes = ShapeInferencer.infer_shapes(model)
+"""
+
+import logging
+
+import onnx
+import sympy
+from onnx import helper
+
+from .context import InferenceContext
+from .registry import get_all_aten_handlers, get_all_shape_handlers, get_aten_handler, get_shape_handler
+from .utils import (
+    get_attribute,
+    get_opset,
+    get_shape_from_type_proto,
+    get_shape_from_value_info,
+    is_literal,
+    is_sequence,
+)
+
+# Import all handlers to trigger registration
+from . import aten_ops  # noqa: F401
+from . import contrib_ops  # noqa: F401
+from . import standard_ops  # noqa: F401
+
+logger = logging.getLogger(__name__)
+
+
+class ShapeInferencer:
+    """Main class for performing symbolic shape inference on ONNX models."""
+
+    def __init__(self, int_max=2**31 - 1, auto_merge=False, guess_output_rank=False, verbose=0, prefix=""):
+        """Initialize the ShapeInferencer.
+
+        Args:
+            int_max: Maximum value for unbounded integers.
+            auto_merge: Whether to automatically merge conflicting dimensions.
+            guess_output_rank: Whether to guess output rank from input.
+            verbose: Logging verbosity level.
+            prefix: Prefix for generated symbolic dimension names.
+        """
+        self.int_max_ = int_max
+        self.auto_merge_ = auto_merge
+        self.guess_output_rank_ = guess_output_rank
+        self.verbose_ = verbose
+        self.prefix_ = prefix
+
+    def _infer_impl(self, ctx, start_sympy_data=None):
+        """Main inference implementation loop."""
+        ctx.sympy_data_ = start_sympy_data or {}
+        ctx.apply_suggested_merge(graph_input_only=True)
+        ctx.input_symbols_ = set()
+
+        # Process graph inputs
+        for i in ctx.out_mp_.graph.input:
+            input_shape = get_shape_from_value_info(i)
+            if input_shape is None:
+                continue
+
+            if is_sequence(i.type):
+                input_dims = i.type.sequence_type.elem_type.tensor_type.shape.dim
+            else:
+                input_dims = i.type.tensor_type.shape.dim
+
+            for i_dim, dim in enumerate(input_shape):
+                if dim is None:
+                    input_dims[i_dim].dim_param = str(ctx.new_symbolic_dim(i.name, i_dim))
+
+            ctx.input_symbols_.update([d for d in input_shape if type(d) == str])
+
+        for s in ctx.input_symbols_:
+            if s in ctx.suggested_merge_:
+                s_merge = ctx.suggested_merge_[s]
+                assert s_merge in ctx.symbolic_dims_
+                ctx.symbolic_dims_[s] = ctx.symbolic_dims_[s_merge]
+            else:
+                ctx.symbolic_dims_[s] = sympy.Symbol(s, integer=True, positive=True)
+
+        # Compute prerequisite for node for topological sort
+        prereq_for_node = {}
+
+        def get_prereq(node):
+            names = {i for i in node.input if i}
+            subgraphs = []
+            if node.op_type == "If":
+                subgraphs = [get_attribute(node, "then_branch"), get_attribute(node, "else_branch")]
+            elif node.op_type in {"Loop", "Scan"}:
+                subgraphs = [get_attribute(node, "body")]
+            for g in subgraphs:
+                g_outputs_and_initializers = {i.name for i in g.initializer}
+                g_prereq = set()
+                for n in g.node:
+                    g_outputs_and_initializers.update(n.output)
+                for n in g.node:
+                    g_prereq.update([i for i in get_prereq(n) if i not in g_outputs_and_initializers])
+                names.update(g_prereq)
+                for i in g.input:
+                    if i.name in names:
+                        names.remove(i.name)
+            return names
+
+        for n in ctx.out_mp_.graph.node:
+            prereq_for_node[n.output[0]] = get_prereq(n)
+
+        # Topological sort nodes
+        sorted_nodes = []
+        sorted_known_vi = {i.name for i in list(ctx.out_mp_.graph.input) + list(ctx.out_mp_.graph.initializer)}
+        if any(o.name in sorted_known_vi for o in ctx.out_mp_.graph.output):
+            sorted_nodes = ctx.out_mp_.graph.node
+        else:
+            while any(o.name not in sorted_known_vi for o in ctx.out_mp_.graph.output):
+                old_sorted_nodes_len = len(sorted_nodes)
+                for node in ctx.out_mp_.graph.node:
+                    if node.output[0] not in sorted_known_vi and all(
+                        i in sorted_known_vi for i in prereq_for_node[node.output[0]] if i
+                    ):
+                        sorted_known_vi.update(node.output)
+                        sorted_nodes.append(node)
+                if old_sorted_nodes_len == len(sorted_nodes) and not all(
+                    o.name in sorted_known_vi for o in ctx.out_mp_.graph.output
+                ):
+                    raise Exception("Invalid model with cyclic graph")
+
+        # Get handlers
+        shape_handlers = get_all_shape_handlers()
+        aten_handlers = get_all_aten_handlers()
+
+        # Process each node
+        for node in sorted_nodes:
+            assert all([i in ctx.known_vi_ for i in node.input if i])
+            ctx.onnx_infer_single_node(node)
+            known_aten_op = False
+
+            # Try standard handlers first
+            handler = get_shape_handler(node.op_type)
+            if handler is not None:
+                handler.infer_shape(node, ctx)
+            elif node.op_type == "ConvTranspose":
+                vi = ctx.known_vi_[node.output[0]]
+                if len(vi.type.tensor_type.shape.dim) == 0:
+                    vi.type.tensor_type.elem_type = onnx.TensorProto.UNDEFINED
+            elif node.op_type == "ATen" and node.domain == "org.pytorch.aten":
+                for attr in node.attribute:
+                    if attr.name == "operator":
+                        aten_op_name = attr.s.decode("utf-8") if isinstance(attr.s, bytes) else attr.s
+                        aten_handler = get_aten_handler(aten_op_name)
+                        if aten_handler is not None:
+                            known_aten_op = True
+                            aten_handler.infer_shape(node, ctx)
+                        break
+
+            if ctx.verbose_ > 2:
+                logger.debug(node.op_type + ": " + node.name)
+                for i, name in enumerate(node.input):
+                    logger.debug(f"  Input {i}: {name} {'initializer' if name in ctx.initializers_ else ''}")
+
+            # Handle dimension merging for broadcast ops
+            if node.op_type in {
+                "Add",
+                "Sub",
+                "Mul",
+                "Div",
+                "MatMul",
+                "MatMulInteger",
+                "MatMulInteger16",
+                "Where",
+                "Sum",
+            }:
+                vi = ctx.known_vi_[node.output[0]]
+                out_rank = len(get_shape_from_type_proto(vi.type))
+                in_shapes = [ctx.get_shape(node, i) for i in range(len(node.input))]
+                for d in range(out_rank - (2 if node.op_type in {"MatMul", "MatMulInteger", "MatMulInteger16"} else 0)):
+                    in_dims = [s[len(s) - out_rank + d] for s in in_shapes if len(s) + d >= out_rank]
+                    if len(in_dims) > 1:
+                        ctx.check_merged_dims(in_dims, allow_broadcast=True)
+
+            # Process outputs
+            for i_o in range(len(node.output)):
+                if node.op_type in {"SkipLayerNormalization", "SkipSimplifiedLayerNormalization"} and i_o in {1, 2}:
+                    continue
+                if node.op_type == "RotaryEmbedding" and len(node.output) > 1:
+                    continue
+
+                vi = ctx.known_vi_[node.output[i_o]]
+                out_type = vi.type
+                out_type_kind = out_type.WhichOneof("value")
+
+                if out_type_kind not in {"tensor_type", "sparse_tensor_type", None}:
+                    if ctx.verbose_ > 2:
+                        if out_type_kind == "sequence_type":
+                            seq_cls_type = out_type.sequence_type.elem_type.WhichOneof("value")
+                            if seq_cls_type == "tensor_type":
+                                logger.debug(
+                                    f"  {node.output[i_o]}: sequence of {str(get_shape_from_value_info(vi))} "
+                                    f"{onnx.TensorProto.DataType.Name(vi.type.sequence_type.elem_type.tensor_type.elem_type)}"
+                                )
+                            else:
+                                logger.debug(f"  {node.output[i_o]}: sequence of {seq_cls_type}")
+                        else:
+                            logger.debug(f"  {node.output[i_o]}: {out_type_kind}")
+                    continue
+
+                out_shape = get_shape_from_value_info(vi)
+                out_type_undefined = out_type.tensor_type.elem_type == onnx.TensorProto.UNDEFINED
+                if ctx.verbose_ > 2:
+                    logger.debug(
+                        f"  {node.output[i_o]}: {out_shape!s} {onnx.TensorProto.DataType.Name(vi.type.tensor_type.elem_type)}"
+                    )
+                    if node.output[i_o] in ctx.sympy_data_:
+                        logger.debug("  Sympy Data: " + str(ctx.sympy_data_[node.output[i_o]]))
+
+                if (out_shape is not None and (None in out_shape or ctx.is_shape_contains_none_dim(out_shape))) or out_type_undefined:
+                    if ctx.auto_merge_:
+                        if node.op_type in {
+                            "Add",
+                            "Sub",
+                            "Mul",
+                            "Div",
+                            "MatMul",
+                            "MatMulInteger",
+                            "MatMulInteger16",
+                            "Concat",
+                            "Where",
+                            "Sum",
+                            "Equal",
+                            "Less",
+                            "Greater",
+                            "LessOrEqual",
+                            "GreaterOrEqual",
+                            "Min",
+                            "Max",
+                        }:
+                            shapes = [ctx.get_shape(node, i) for i in range(len(node.input))]
+                            if node.op_type in {"MatMul", "MatMulInteger", "MatMulInteger16"} and (
+                                None in out_shape or ctx.is_shape_contains_none_dim(out_shape)
+                            ):
+                                if None in out_shape:
+                                    idx = out_shape.index(None)
+                                else:
+                                    idx = out_shape.index(ctx.is_shape_contains_none_dim(out_shape))
+                                dim_idx = [len(s) - len(out_shape) + idx for s in shapes]
+                                assert len(shapes[0]) > 2 and dim_idx[0] < len(shapes[0]) - 2
+                                assert len(shapes[1]) > 2 and dim_idx[1] < len(shapes[1]) - 2
+                        elif node.op_type == "Expand":
+                            shapes = [ctx.get_shape(node, 0), ctx.get_value(node, 1)]
+                        else:
+                            shapes = []
+
+                        if shapes:
+                            for idx in range(len(out_shape)):
+                                if out_shape[idx] is not None and not ctx.is_none_dim(out_shape[idx]):
+                                    continue
+                                dim_idx = [len(s) - len(out_shape) + idx for s in shapes]
+                                if dim_idx:
+                                    ctx.add_suggested_merge(
+                                        [s[i] if is_literal(s[i]) else str(s[i]) for s, i in zip(shapes, dim_idx) if i >= 0]
+                                    )
+                            ctx.run_ = True
+                        else:
+                            ctx.run_ = False
+                    else:
+                        ctx.run_ = False
+
+                    if not ctx.run_ and handler is None and not known_aten_op:
+                        is_unknown_op = out_type_undefined and (out_shape is None or len(out_shape) == 0)
+                        if is_unknown_op:
+                            out_rank = ctx.get_shape_rank(node, 0) if ctx.guess_output_rank_ else -1
+                        else:
+                            out_rank = len(out_shape)
+
+                        if out_rank >= 0:
+                            new_shape = ctx.new_symbolic_shape(out_rank, node, i_o)
+                            if out_type_undefined:
+                                out_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+                            else:
+                                out_dtype = vi.type.tensor_type.elem_type
+                            from .utils import get_shape_from_sympy_shape
+
+                            vi.CopyFrom(
+                                helper.make_tensor_value_info(vi.name, out_dtype, get_shape_from_sympy_shape(new_shape))
+                            )
+
+                            if ctx.verbose_ > 0:
+                                if is_unknown_op:
+                                    logger.debug(f"Possible unknown op: {node.op_type} node: {node.name}, guessing {vi.name} shape")
+                                if ctx.verbose_ > 2:
+                                    logger.debug(f"  {node.output[i_o]}: {new_shape!s} {vi.type.tensor_type.elem_type}")
+                            ctx.run_ = True
+                            continue
+
+                    if ctx.verbose_ > 0 or not ctx.auto_merge_ or out_type_undefined:
+                        logger.debug("Stopping at incomplete shape inference at " + node.op_type + ": " + node.name)
+                        logger.debug("node inputs:")
+                        for i in node.input:
+                            if i in ctx.known_vi_:
+                                logger.debug(ctx.known_vi_[i])
+                            else:
+                                logger.debug(f"not in known_vi_ for {i}")
+                        logger.debug("node outputs:")
+                        for o in node.output:
+                            if o in ctx.known_vi_:
+                                logger.debug(ctx.known_vi_[o])
+                            else:
+                                logger.debug(f"not in known_vi_ for {o}")
+                        if ctx.auto_merge_ and not out_type_undefined:
+                            logger.debug("Merging: " + str(ctx.suggested_merge_))
+                    return False
+
+        ctx.run_ = False
+        return True
+
+    def _update_output_from_vi(self, ctx):
+        """Update output attributes using known value information dictionary."""
+        for output in ctx.out_mp_.graph.output:
+            if output.name in ctx.known_vi_:
+                output.CopyFrom(ctx.known_vi_[output.name])
+
+    @staticmethod
+    def infer_shapes(in_mp, int_max=2**31 - 1, auto_merge=False, guess_output_rank=False, verbose=0):
+        """Perform symbolic shape inference on an ONNX model.
+
+        Args:
+            in_mp: The input ONNX ModelProto.
+            int_max: Maximum value for unbounded integers.
+            auto_merge: Whether to automatically merge conflicting dimensions.
+            guess_output_rank: Whether to guess output rank from input.
+            verbose: Logging verbosity level.
+
+        Returns:
+            The model with inferred shapes.
+
+        Raises:
+            Exception: If shape inference is incomplete.
+        """
+        onnx_opset = get_opset(in_mp)
+        if (not onnx_opset) or onnx_opset < 7:
+            logger.warning("Only support models of onnx opset 7 and above.")
+            return None
+
+        inferencer = ShapeInferencer(int_max, auto_merge, guess_output_rank, verbose)
+
+        # Create inference context
+        ctx = InferenceContext(
+            in_mp,
+            int_max=int_max,
+            auto_merge=auto_merge,
+            guess_output_rank=guess_output_rank,
+            verbose=verbose,
+        )
+        ctx.preprocess()
+
+        all_shapes_inferred = False
+        while ctx.run_:
+            all_shapes_inferred = inferencer._infer_impl(ctx)
+
+        inferencer._update_output_from_vi(ctx)
+
+        if not all_shapes_inferred:
+            raise Exception("Incomplete symbolic shape inference")
+
+        return ctx.out_mp_
+
+
+# For backward compatibility
+SymbolicShapeInference = ShapeInferencer

onnxslim/core/shape_inference/aten_ops/__init__.py

@@ -0,0 +1,16 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""PyTorch ATen operator shape handlers."""
+
+from . import bitwise_or
+from . import diagonal
+from . import pool2d
+from . import min_max
+from . import multinomial
+from . import unfold
+from . import argmax
+from . import group_norm
+from . import upsample
+from . import embedding
+from . import numpy_t