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

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

@@ -0,0 +1,119 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Einsum operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class EinsumHandler(ShapeHandler):
+    """Handler for Einsum operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Einsum"
+
+    def infer_shape(self, node, ctx) -> None:
+        equation = get_attribute(node, "equation")
+        equation = equation.replace(b" ", b"")
+        mid_index = equation.find(b"->")
+        left_equation = equation[:mid_index] if mid_index != -1 else equation
+
+        num_operands = 0
+        num_ellipsis = 0
+        num_ellipsis_indices = 0
+        num_labels = 0
+        ellipsis_flag = True
+        dims_value = []
+        ellipsis_dims_value = []
+
+        label_maps = {}
+        repeated_labels = set()
+
+        terms = left_equation.split(b",")
+        for term in terms:
+            ellipsis_index = term.find(b"...")
+            shape = ctx.get_shape(node, num_operands)
+            rank = len(shape)
+            ellipsis_dims = 0
+            term_size = 0
+            num_illegal_char = 0
+
+            for i in range(len(term)):
+                if term[i] != 46:
+                    term_size = term_size + 1
+
+            index = 0
+            while index < len(term):
+                if index == ellipsis_index:
+                    ellipsis_dims = rank - term_size
+                    if ellipsis_flag:
+                        ellipsis_flag = False
+                        for i in range(ellipsis_dims):
+                            ellipsis_dims_value.append(shape[index + i - num_illegal_char])
+                    else:
+                        for i in range(ellipsis_dims):
+                            shape_dim = shape[index + i - num_illegal_char]
+                            current_dim = ellipsis_dims_value[i]
+                            ellipsis_dims_value[i] = max(current_dim, shape_dim)
+
+                    num_illegal_char += 3
+                    index += 3
+                    continue
+
+                elif term[index] == 46:
+                    num_illegal_char += 1
+                    index += 1
+                    continue
+
+                char = term[index]
+                if char not in label_maps:
+                    label_maps[char] = num_labels
+                    dims_value.append(shape[index + ellipsis_dims - num_illegal_char])
+                    num_labels += 1
+                else:
+                    repeated_labels.add(char)
+
+                index += 1
+
+            if ellipsis_index != -1:
+                if num_ellipsis == 0:
+                    if rank < term_size:
+                        raise ValueError("Ellipsis represents incompatible dimensions.")
+                    num_ellipsis_indices = rank - term_size
+                else:
+                    if num_ellipsis_indices != rank - term_size:
+                        raise ValueError("Ellipsis represents incompatible dimensions.")
+                num_ellipsis += 1
+            else:
+                if rank != term_size:
+                    raise ValueError("Rank of input ", num_operands, " does not match the equation indices.")
+            num_operands += 1
+
+        new_sympy_shape = []
+        if mid_index != -1:
+            right_equation = equation[mid_index + 2 :]
+            right_ellipsis_index = right_equation.find(b"...")
+            if right_ellipsis_index != -1:
+                for i in range(num_ellipsis_indices):
+                    new_sympy_shape.append(ellipsis_dims_value[i])
+            for c in right_equation:
+                if c != 46:
+                    new_sympy_shape.append(dims_value[label_maps[c]])
+        else:
+            for i in range(num_ellipsis_indices):
+                new_sympy_shape.append(ellipsis_dims_value[i])
+            for label, idx in label_maps.items():
+                if label not in repeated_labels:
+                    new_sympy_shape.append(dims_value[idx])
+
+        output_dtype = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, new_sympy_shape))
+
+
+register_shape_handler(EinsumHandler())

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ReduceProd operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, sympy_reduce_product
+
+
+class ReduceProdHandler(ShapeHandler):
+    """Handler for ReduceProd operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ReduceProd"
+
+    def infer_shape(self, node, ctx) -> None:
+        axes = get_attribute(node, "axes")
+        keep_dims = get_attribute(node, "keepdims", 1)
+        if keep_dims == 0 and axes == [0]:
+            data = ctx.get_int_or_float_values(node)[0]
+            if data is not None:
+                ctx.sympy_data_[node.output[0]] = sympy_reduce_product(data)
+
+
+register_shape_handler(ReduceProdHandler())

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

@@ -0,0 +1,53 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ReduceSum operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_opset, get_shape_from_sympy_shape, handle_negative_axis
+
+
+class ReduceSumHandler(ShapeHandler):
+    """Handler for ReduceSum operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ReduceSum"
+
+    def infer_shape(self, node, ctx) -> None:
+        keep_dims = get_attribute(node, "keepdims", 1)
+        if get_opset(ctx.out_mp_) >= 13 and len(node.input) > 1:
+            axes = ctx.try_get_value(node, 1)
+            vi = ctx.known_vi_[node.output[0]]
+            if axes is None:
+                assert keep_dims
+                vi.CopyFrom(
+                    helper.make_tensor_value_info(
+                        node.output[0],
+                        ctx.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                        get_shape_from_sympy_shape(ctx.new_symbolic_shape(ctx.get_shape_rank(node, 0), node)),
+                    )
+                )
+            else:
+                shape = ctx.get_shape(node, 0)
+                output_shape = []
+                axes = [handle_negative_axis(a, len(shape)) for a in axes]
+                for i, d in enumerate(shape):
+                    if i in axes:
+                        if keep_dims:
+                            output_shape.append(1)
+                    else:
+                        output_shape.append(d)
+                vi.CopyFrom(
+                    helper.make_tensor_value_info(
+                        node.output[0],
+                        ctx.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                        output_shape,
+                    )
+                )
+
+
+register_shape_handler(ReduceSumHandler())

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

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

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

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

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

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Miscellaneous operator shape handlers."""
+
+from . import constant
+from . import constant_of_shape
+from . import cast
+from . import category_mapper
+from . import compress
+from . import one_hot
+from . import non_max_suppression
+from . import non_zero
+from . import top_k
+from . import range
+from . import resize
+from . import scatter_elements
+from . import array_feature_extractor
+from . import softmax_cross_entropy_loss
+from . import dequantize_linear
+from . import quantize_linear
+from . import relative_position_bias

onnxslim/core/shape_inference/standard_ops/misc/array_feature_extractor.py

@@ -0,0 +1,32 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ArrayFeatureExtractor operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class ArrayFeatureExtractorHandler(ShapeHandler):
+    """Handler for ArrayFeatureExtractor operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ArrayFeatureExtractor"
+
+    def infer_shape(self, node, ctx) -> None:
+        data_shape = ctx.get_shape(node, 0)
+        indices_shape = ctx.get_shape(node, 1)
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(
+            helper.make_tensor_value_info(
+                node.output[0],
+                ctx.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                data_shape[:-1] + indices_shape,
+            )
+        )
+
+
+register_shape_handler(ArrayFeatureExtractorHandler())

onnxslim/core/shape_inference/standard_ops/misc/cast.py

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

onnxslim/core/shape_inference/standard_ops/misc/category_mapper.py

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for CategoryMapper operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class CategoryMapperHandler(ShapeHandler):
+    """Handler for CategoryMapper operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "CategoryMapper"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_type = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+        if input_type == onnx.TensorProto.STRING:
+            output_type = onnx.TensorProto.INT64
+        else:
+            output_type = onnx.TensorProto.STRING
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_type, ctx.get_shape(node, 0)))
+
+
+register_shape_handler(CategoryMapperHandler())

onnxslim/core/shape_inference/standard_ops/misc/compress.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Compress operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, handle_negative_axis
+
+
+class CompressHandler(ShapeHandler):
+    """Handler for Compress operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Compress"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_shape = ctx.get_shape(node, 0)
+        compress_len = str(ctx.new_symbolic_dim_from_output(node))
+        axis = get_attribute(node, "axis")
+        if axis is None:
+            output_shape = [compress_len]
+        else:
+            output_shape = input_shape
+            output_shape[handle_negative_axis(axis, len(input_shape))] = compress_len
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(
+            helper.make_tensor_value_info(
+                node.output[0],
+                ctx.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                output_shape,
+            )
+        )
+
+
+register_shape_handler(CompressHandler())

onnxslim/core/shape_inference/standard_ops/misc/constant.py

@@ -0,0 +1,27 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Constant operator."""
+
+from onnx import numpy_helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class ConstantHandler(ShapeHandler):
+    """Handler for Constant operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Constant"
+
+    def infer_shape(self, node, ctx) -> None:
+        t = get_attribute(node, "value")
+        t.name = node.output[0]
+        ctx.initializers_[node.output[0]] = t
+        ctx.sympy_data_[node.output[0]] = numpy_helper.to_array(t)
+
+
+register_shape_handler(ConstantHandler())

onnxslim/core/shape_inference/standard_ops/misc/constant_of_shape.py

@@ -0,0 +1,45 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ConstantOfShape operator."""
+
+import numpy as np
+import onnx
+from onnx import helper, numpy_helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_shape_from_sympy_shape, is_literal
+
+
+class ConstantOfShapeHandler(ShapeHandler):
+    """Handler for ConstantOfShape operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ConstantOfShape"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_int_or_float_values(node)[0]
+        vi = ctx.known_vi_[node.output[0]]
+        if sympy_shape is not None:
+            if type(sympy_shape) != list:
+                sympy_shape = [sympy_shape]
+            ctx.update_computed_dims(sympy_shape)
+            if vi.type.tensor_type.elem_type == onnx.TensorProto.INT64 and all(is_literal(x) for x in sympy_shape):
+                ctx.sympy_data_[node.output[0]] = np.ones(
+                    [int(x) for x in sympy_shape], dtype=np.int64
+                ) * numpy_helper.to_array(get_attribute(node, "value", 0))
+        else:
+            sympy_shape = ctx.new_symbolic_shape(ctx.get_shape(node, 0)[0], node)
+
+        vi.CopyFrom(
+            helper.make_tensor_value_info(
+                node.output[0],
+                vi.type.tensor_type.elem_type,
+                get_shape_from_sympy_shape(sympy_shape),
+            )
+        )
+
+
+register_shape_handler(ConstantOfShapeHandler())

onnxslim/core/shape_inference/standard_ops/misc/dequantize_linear.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for DequantizeLinear operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class DequantizeLinearHandler(ShapeHandler):
+    """Handler for DequantizeLinear operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "DequantizeLinear"
+
+    def infer_shape(self, node, ctx) -> None:
+        output_dtype = ctx.known_vi_[node.input[1]].type.tensor_type.elem_type
+        output_shape = ctx.get_shape(node, 0)
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], output_dtype, output_shape))
+
+
+register_shape_handler(DequantizeLinearHandler())

onnxslim/core/shape_inference/standard_ops/misc/non_max_suppression.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for NonMaxSuppression operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class NonMaxSuppressionHandler(ShapeHandler):
+    """Handler for NonMaxSuppression operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "NonMaxSuppression"
+
+    def infer_shape(self, node, ctx) -> None:
+        selected = str(ctx.new_symbolic_dim_from_output(node))
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64, [selected, 3]))
+
+
+register_shape_handler(NonMaxSuppressionHandler())

onnxslim/core/shape_inference/standard_ops/misc/non_zero.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for NonZero operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class NonZeroHandler(ShapeHandler):
+    """Handler for NonZero operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "NonZero"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_rank = ctx.get_shape_rank(node, 0)
+        nz_len = str(ctx.new_symbolic_dim_from_output(node, 0, 1))
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(helper.make_tensor_value_info(node.output[0], vi.type.tensor_type.elem_type, [input_rank, nz_len]))
+
+
+register_shape_handler(NonZeroHandler())