onnxslim 0.1.82__py3-none-any.whl → 0.1.84__py3-none-any.whl

onnxslim/core/shape_inference/contrib_ops/misc/quick_gelu.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for QuickGelu operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class QuickGeluHandler(ShapeHandler):
+    """Handler for QuickGelu operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "QuickGelu"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(QuickGeluHandler())

onnxslim/core/shape_inference/contrib_ops/misc/rotary_embedding.py

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for RotaryEmbedding operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class RotaryEmbeddingHandler(ShapeHandler):
+    """Handler for RotaryEmbedding operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "RotaryEmbedding"
+
+    def infer_shape(self, node, ctx) -> None:
+        if len(node.output) == 1:
+            ctx.propagate_shape_and_type(node)
+        elif len(node.output) == 2:
+            # Extraneous constant nodes outputted by RotaryEmbedding function made with `export_modules_as_functions`
+            ctx.propagate_shape_and_type(node, input_index=1, output_index=0)
+            ctx.propagate_shape_and_type(node, input_index=0, output_index=1)
+        elif len(node.output) == 3:
+            # Extraneous constant nodes outputted by RotaryEmbedding function made with `export_modules_as_functions`
+            ctx.propagate_shape_and_type(node, input_index=1, output_index=0)
+            ctx.propagate_shape_and_type(node, input_index=1, output_index=1)
+            ctx.propagate_shape_and_type(node, input_index=0, output_index=2)
+
+
+register_shape_handler(RotaryEmbeddingHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/__init__.py

@@ -0,0 +1,12 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Normalization-related contrib operator shape handlers."""
+
+from . import layer_normalization
+from . import simplified_layer_normalization
+from . import skip_layer_normalization
+from . import skip_simplified_layer_normalization
+from . import group_norm
+from . import skip_group_norm
+from . import embed_layer_normalization

onnxslim/core/shape_inference/contrib_ops/normalization/embed_layer_normalization.py

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for EmbedLayerNormalization operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class EmbedLayerNormalizationHandler(ShapeHandler):
+    """Handler for EmbedLayerNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "EmbedLayerNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_ids_shape = ctx.get_shape(node, 0)
+        word_embedding_shape = ctx.get_shape(node, 2)
+        assert len(input_ids_shape) == 2 and len(word_embedding_shape) == 2
+        output_shape = [*input_ids_shape, word_embedding_shape[1]]
+
+        word_embedding_dtype = ctx.known_vi_[node.input[2]].type.tensor_type.elem_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], word_embedding_dtype, output_shape))
+
+        if len(node.output) > 1 and node.output[1]:
+            mask_index_shape = [input_ids_shape[0]]
+            vi = ctx.known_vi_[node.output[1]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[1], onnx.TensorProto.INT32, mask_index_shape))
+
+        if len(node.output) > 2:
+            # Optional output of add before layer normalization is done
+            vi = ctx.known_vi_[node.output[2]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[2], word_embedding_dtype, output_shape))
+
+
+register_shape_handler(EmbedLayerNormalizationHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/group_norm.py

@@ -0,0 +1,21 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for GroupNorm operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class GroupNormHandler(ShapeHandler):
+    """Handler for GroupNorm operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "GroupNorm"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+
+register_shape_handler(GroupNormHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/layer_normalization.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for LayerNormalization operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, handle_negative_axis
+
+
+class LayerNormalizationHandler(ShapeHandler):
+    """Handler for LayerNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "LayerNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+        if len(node.output) > 1:
+            axis = get_attribute(node, "axis")
+            if axis is None:
+                axis = -1
+            x_shape = ctx.get_shape(node, 0)
+            if x_shape is not None:
+                rank = len(x_shape)
+                axis = handle_negative_axis(axis, rank)
+                mean_shape = x_shape[:axis] + [1 for _ in range(rank - axis)]
+                mean_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+                if mean_dtype in {onnx.TensorProto.FLOAT16, onnx.TensorProto.BFLOAT16}:
+                    mean_dtype = onnx.TensorProto.FLOAT
+                vi = ctx.known_vi_[node.output[1]]
+                vi.CopyFrom(helper.make_tensor_value_info(node.output[1], mean_dtype, mean_shape))
+                if len(node.output) > 2:
+                    vi = ctx.known_vi_[node.output[2]]
+                    vi.CopyFrom(helper.make_tensor_value_info(node.output[2], mean_dtype, mean_shape))
+
+
+register_shape_handler(LayerNormalizationHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/simplified_layer_normalization.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SimplifiedLayerNormalization operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from .layer_normalization import LayerNormalizationHandler
+
+
+class SimplifiedLayerNormalizationHandler(ShapeHandler):
+    """Handler for SimplifiedLayerNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SimplifiedLayerNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        # Reuse LayerNormalization handler
+        LayerNormalizationHandler().infer_shape(node, ctx)
+
+
+register_shape_handler(SimplifiedLayerNormalizationHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/skip_group_norm.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SkipGroupNorm operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class SkipGroupNormHandler(ShapeHandler):
+    """Handler for SkipGroupNorm operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SkipGroupNorm"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node, 0, 0)
+        if len(node.output) > 1:
+            ctx.propagate_shape_and_type(node, 0, 1)
+
+
+register_shape_handler(SkipGroupNormHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/skip_layer_normalization.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SkipLayerNormalization operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class SkipLayerNormalizationHandler(ShapeHandler):
+    """Handler for SkipLayerNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SkipLayerNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+        # If the SkipLayerNormalization node contains the optional
+        # output for inference, infer the shape and type for it too
+        if len(node.output) > 3:
+            ctx.propagate_shape_and_type(node, 0, 3)
+
+
+register_shape_handler(SkipLayerNormalizationHandler())

onnxslim/core/shape_inference/contrib_ops/normalization/skip_simplified_layer_normalization.py

@@ -0,0 +1,23 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SkipSimplifiedLayerNormalization operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from .skip_layer_normalization import SkipLayerNormalizationHandler
+
+
+class SkipSimplifiedLayerNormalizationHandler(ShapeHandler):
+    """Handler for SkipSimplifiedLayerNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SkipSimplifiedLayerNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        # Reuse SkipLayerNormalization handler
+        SkipLayerNormalizationHandler().infer_shape(node, ctx)
+
+
+register_shape_handler(SkipSimplifiedLayerNormalizationHandler())

onnxslim/core/shape_inference/registry.py

@@ -0,0 +1,90 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Registry for shape inference handlers."""
+
+from collections import OrderedDict
+
+# Global registries for shape handlers
+SHAPE_HANDLERS = OrderedDict()
+ATEN_SHAPE_HANDLERS = OrderedDict()
+
+
+def register_shape_handler(handler):
+    """Register a shape handler for a specific ONNX operator type.
+
+    Args:
+        handler: A ShapeHandler instance to register.
+
+    Returns:
+        The registered handler.
+
+    Raises:
+        ValueError: If a handler for the same op_type is already registered.
+    """
+    op_type = handler.op_type
+    if op_type in SHAPE_HANDLERS:
+        raise ValueError(f"Handler for op_type '{op_type}' is already registered")
+    SHAPE_HANDLERS[op_type] = handler
+    return handler
+
+
+def register_aten_handler(handler):
+    """Register a shape handler for a PyTorch ATen operator.
+
+    Args:
+        handler: A ShapeHandler instance to register.
+
+    Returns:
+        The registered handler.
+
+    Raises:
+        ValueError: If a handler for the same op_name is already registered.
+    """
+    op_name = handler.op_type
+    if op_name in ATEN_SHAPE_HANDLERS:
+        raise ValueError(f"Handler for ATen op '{op_name}' is already registered")
+    ATEN_SHAPE_HANDLERS[op_name] = handler
+    return handler
+
+
+def get_shape_handler(op_type):
+    """Get the shape handler for a given ONNX operator type.
+
+    Args:
+        op_type: The ONNX operator type string.
+
+    Returns:
+        The registered ShapeHandler or None if not found.
+    """
+    return SHAPE_HANDLERS.get(op_type)
+
+
+def get_aten_handler(op_name):
+    """Get the shape handler for a given ATen operator name.
+
+    Args:
+        op_name: The ATen operator name string.
+
+    Returns:
+        The registered ShapeHandler or None if not found.
+    """
+    return ATEN_SHAPE_HANDLERS.get(op_name)
+
+
+def get_all_shape_handlers():
+    """Get all registered shape handlers.
+
+    Returns:
+        OrderedDict of all registered shape handlers.
+    """
+    return SHAPE_HANDLERS.copy()
+
+
+def get_all_aten_handlers():
+    """Get all registered ATen shape handlers.
+
+    Returns:
+        OrderedDict of all registered ATen shape handlers.
+    """
+    return ATEN_SHAPE_HANDLERS.copy()

onnxslim/core/shape_inference/standard_ops/__init__.py

@@ -0,0 +1,11 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Standard ONNX operator shape handlers."""
+
+from . import tensor
+from . import math
+from . import nn
+from . import control_flow
+from . import sequence
+from . import misc

onnxslim/core/shape_inference/standard_ops/control_flow/__init__.py

@@ -0,0 +1,8 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Control flow operator shape handlers."""
+
+from . import if_op
+from . import loop
+from . import scan

onnxslim/core/shape_inference/standard_ops/control_flow/if_op.py

@@ -0,0 +1,43 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for If operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import as_scalar, get_attribute
+
+
+class IfHandler(ShapeHandler):
+    """Handler for If operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "If"
+
+    def infer_shape(self, node, ctx) -> None:
+        subgraphs = [
+            get_attribute(node, "then_branch"),
+            get_attribute(node, "else_branch"),
+        ]
+
+        cond = ctx.try_get_value(node, 0)
+
+        for i_sub, subgraph in enumerate(subgraphs):
+            subgraph_infer = ctx.onnx_infer_subgraph(node, subgraph, use_node_input=False)
+            for i_out in range(len(node.output)):
+                vi = ctx.known_vi_[node.output[i_out]]
+                if i_sub == 0:
+                    vi.CopyFrom(subgraph.output[i_out])
+                    vi.name = node.output[i_out]
+                else:
+                    ctx.fuse_tensor_type(node, i_out, vi.type, subgraph.output[i_out].type)
+                if (
+                    cond is not None
+                    and i_sub == (0 if as_scalar(cond) > 0 else 1)
+                    and subgraph.output[i_out].name in subgraph_infer.sympy_data_
+                ):
+                    ctx.sympy_data_[vi.name] = subgraph_infer.sympy_data_[subgraph.output[i_out].name]
+
+
+register_shape_handler(IfHandler())

onnxslim/core/shape_inference/standard_ops/control_flow/loop.py

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Loop operator."""
+
+import logging
+
+import onnx
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_shape_from_value_info, is_sequence
+
+logger = logging.getLogger(__name__)
+
+
+class LoopHandler(ShapeHandler):
+    """Handler for Loop operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Loop"
+
+    def infer_shape(self, node, ctx) -> None:
+        subgraph = get_attribute(node, "body")
+        assert len(subgraph.input) == len(node.input)
+        num_loop_carried = len(node.input) - 2
+
+        for i, si in enumerate(subgraph.input):
+            si_name = si.name
+            si.CopyFrom(ctx.known_vi_[node.input[i]])
+            si.name = si_name
+
+        ctx.onnx_infer_subgraph(node, subgraph)
+
+        need_second_infer = False
+        for i_out in range(1, num_loop_carried + 1):
+            so = subgraph.output[i_out]
+            so_shape = get_shape_from_value_info(so)
+            if is_sequence(so.type):
+                if so_shape and None in so_shape:
+                    subgraph.input[i_out + 1].type.sequence_type.elem_type.CopyFrom(so.type.sequence_type.elem_type)
+                    need_second_infer = True
+            else:
+                si = subgraph.input[i_out + 1]
+                si_shape = get_shape_from_value_info(si)
+                for di, dims in enumerate(zip(si_shape, so_shape)):
+                    if dims[0] != dims[1]:
+                        new_dim = onnx.TensorShapeProto.Dimension()
+                        new_dim.dim_param = str(ctx.new_symbolic_dim_from_output(node, i_out, di))
+                        si.type.tensor_type.shape.dim[di].CopyFrom(new_dim)
+                        so.type.tensor_type.shape.dim[di].CopyFrom(new_dim)
+                        need_second_infer = True
+
+        if need_second_infer:
+            if ctx.verbose_ > 2:
+                logger.debug(f"Rerun Loop: {node.name}({node.output[0]}...), because of sequence in loop carried variables")
+            ctx.onnx_infer_subgraph(node, subgraph, inc_subgraph_id=False)
+
+        loop_iter_dim = str(ctx.new_symbolic_dim_from_output(node))
+        for i in range(len(node.output)):
+            vi = ctx.known_vi_[node.output[i]]
+            vi.CopyFrom(subgraph.output[i + 1])
+            if i >= num_loop_carried:
+                assert not is_sequence(vi.type)
+                subgraph_vi_dim = subgraph.output[i + 1].type.tensor_type.shape.dim
+                vi.type.tensor_type.shape.ClearField("dim")
+                vi_dim = vi.type.tensor_type.shape.dim
+                vi_dim.add().dim_param = loop_iter_dim
+                vi_dim.extend(list(subgraph_vi_dim))
+            vi.name = node.output[i]
+
+
+register_shape_handler(LoopHandler())

onnxslim/core/shape_inference/standard_ops/control_flow/scan.py

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Scan operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_shape_from_type_proto, handle_negative_axis
+
+
+class ScanHandler(ShapeHandler):
+    """Handler for Scan operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Scan"
+
+    def infer_shape(self, node, ctx) -> None:
+        subgraph = get_attribute(node, "body")
+        num_scan_inputs = get_attribute(node, "num_scan_inputs")
+        scan_input_axes = get_attribute(node, "scan_input_axes", [0] * num_scan_inputs)
+        num_scan_states = len(node.input) - num_scan_inputs
+        scan_input_axes = [
+            handle_negative_axis(ax, ctx.get_shape_rank(node, i + num_scan_states)) for i, ax in enumerate(scan_input_axes)
+        ]
+
+        assert len(subgraph.input) >= len(node.input)
+        subgraph_inputs = subgraph.input[: len(node.input)]
+        for i, si in enumerate(subgraph_inputs):
+            subgraph_name = si.name
+            si.CopyFrom(ctx.known_vi_[node.input[i]])
+            if i >= num_scan_states:
+                scan_input_dim = si.type.tensor_type.shape.dim
+                scan_input_dim.remove(scan_input_dim[scan_input_axes[i - num_scan_states]])
+            si.name = subgraph_name
+        ctx.onnx_infer_subgraph(node, subgraph)
+        num_scan_outputs = len(node.output) - num_scan_states
+        scan_output_axes = get_attribute(node, "scan_output_axes", [0] * num_scan_outputs)
+        scan_input_dim = get_shape_from_type_proto(ctx.known_vi_[node.input[-1]].type)[scan_input_axes[-1]]
+        for i, o in enumerate(node.output):
+            vi = ctx.known_vi_[o]
+            if i >= num_scan_states:
+                shape = get_shape_from_type_proto(subgraph.output[i].type)
+                new_dim = handle_negative_axis(scan_output_axes[i - num_scan_states], len(shape) + 1)
+                shape = [*shape[:new_dim], scan_input_dim, *shape[new_dim:]]
+                vi.CopyFrom(helper.make_tensor_value_info(o, subgraph.output[i].type.tensor_type.elem_type, shape))
+            else:
+                vi.CopyFrom(subgraph.output[i])
+            vi.name = o
+
+
+register_shape_handler(ScanHandler())

onnxslim/core/shape_inference/standard_ops/math/__init__.py

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Mathematical operator shape handlers."""
+
+from . import matmul
+from . import matmul_integer
+from . import einsum
+from . import reduce_sum
+from . import reduce_prod
+from . import add
+from . import sub
+from . import mul
+from . import div
+from . import neg
+from . import floor
+from . import min
+from . import max
+from . import equal
+from . import where

onnxslim/core/shape_inference/standard_ops/math/_symbolic_compute.py

@@ -0,0 +1,34 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shared symbolic computation helper for math operators."""
+
+import sympy
+
+from ...utils import is_literal
+
+
+def infer_symbolic_compute_ops(node, ctx):
+    """Handles symbolic computation operations for given node based on predefined functions."""
+    funcs = {
+        "Add": lambda l: l[0] + l[1],
+        "Div": lambda l: (int(l[0] // l[1]) if isinstance(l[0] // l[1], float) else l[0] // l[1]),
+        "Equal": lambda l: l[0] == l[1],
+        "Floor": lambda l: sympy.floor(l[0]),
+        "Max": lambda l: (
+            l[1]
+            if is_literal(l[0]) and int(l[0]) < -ctx.int_max_
+            else (l[0] if is_literal(l[1]) and int(l[1]) < -ctx.int_max_ else sympy.Max(l[0], l[1]))
+        ),
+        "Min": lambda l: (
+            l[1]
+            if is_literal(l[0]) and int(l[0]) > ctx.int_max_
+            else (l[0] if is_literal(l[1]) and int(l[1]) > ctx.int_max_ else sympy.Min(l[0], l[1]))
+        ),
+        "Mul": lambda l: (int(l[0] * l[1]) if isinstance(l[0] * l[1], float) else l[0] * l[1]),
+        "Sub": lambda l: l[0] - l[1],
+        "Where": lambda l: l[1] if l[0] else l[2],
+        "Neg": lambda l: -l[0],
+    }
+    assert node.op_type in funcs
+    ctx.compute_on_sympy_data(node, funcs[node.op_type])

onnxslim/core/shape_inference/standard_ops/math/add.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Add operator."""
+
+from ...base import MultiOpHandler
+from ...registry import register_shape_handler
+from ._symbolic_compute import infer_symbolic_compute_ops
+
+register_shape_handler(MultiOpHandler("Add", infer_symbolic_compute_ops))

onnxslim/core/shape_inference/standard_ops/math/div.py

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Div operator."""
+
+from ...base import MultiOpHandler
+from ...registry import register_shape_handler
+from ._symbolic_compute import infer_symbolic_compute_ops
+
+register_shape_handler(MultiOpHandler("Div", infer_symbolic_compute_ops))