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

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

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for OneHot operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_shape_from_sympy_shape, handle_negative_axis, is_literal
+
+
+class OneHotHandler(ShapeHandler):
+    """Handler for OneHot operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "OneHot"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        depth = ctx.try_get_value(node, 1)
+        axis = get_attribute(node, "axis", -1)
+        axis = handle_negative_axis(axis, len(sympy_shape) + 1)
+        new_shape = get_shape_from_sympy_shape(
+            [
+                *sympy_shape[:axis],
+                depth if is_literal(depth) else ctx.new_symbolic_dim_from_output(node),
+                *sympy_shape[axis:],
+            ]
+        )
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(
+            helper.make_tensor_value_info(
+                node.output[0],
+                ctx.known_vi_[node.input[2]].type.tensor_type.elem_type,
+                new_shape,
+            )
+        )
+
+
+register_shape_handler(OneHotHandler())

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

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for QuantizeLinear operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class QuantizeLinearHandler(ShapeHandler):
+    """Handler for QuantizeLinear operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "QuantizeLinear"
+
+    def infer_shape(self, node, ctx) -> None:
+        output_dtype = onnx.TensorProto.UINT8
+        if len(node.input) > 2 and node.input[2]:
+            output_dtype = ctx.known_vi_[node.input[2]].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(QuantizeLinearHandler())

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

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Range operator."""
+
+import sympy
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import as_scalar, get_shape_from_sympy_shape
+
+
+class RangeHandler(ShapeHandler):
+    """Handler for Range operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Range"
+
+    def infer_shape(self, node, ctx) -> None:
+        vi = ctx.known_vi_[node.output[0]]
+        input_data = ctx.get_int_or_float_values(node)
+        if all(i is not None for i in input_data):
+            start = as_scalar(input_data[0])
+            limit = as_scalar(input_data[1])
+            delta = as_scalar(input_data[2])
+            new_sympy_shape = [sympy.Max(sympy.ceiling((limit - start) / delta), 0)]
+        else:
+            new_sympy_shape = [ctx.new_symbolic_dim_from_output(node)]
+        ctx.update_computed_dims(new_sympy_shape)
+        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(new_sympy_shape),
+            )
+        )
+
+
+register_shape_handler(RangeHandler())

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

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for RelativePositionBias operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class RelativePositionBiasHandler(ShapeHandler):
+    """Handler for RelativePositionBias operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "RelativePositionBias"
+
+    def infer_shape(self, node, ctx) -> None:
+        seq_len = ctx.try_get_value(node, 1)
+        real_seq_len = ctx.try_get_value(node, 2)
+        if seq_len is None or real_seq_len is None:
+            return
+        num_heads = ctx.get_sympy_shape(node, 0)[1]
+        new_shape = [1, num_heads, str(seq_len), str(real_seq_len)]
+        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_shape))
+
+
+register_shape_handler(RelativePositionBiasHandler())

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

@@ -0,0 +1,74 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Resize operator."""
+
+import numpy as np
+import sympy
+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
+
+
+class ResizeHandler(ShapeHandler):
+    """Handler for Resize operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Resize"
+
+    def infer_shape(self, node, ctx) -> None:
+        vi = ctx.known_vi_[node.output[0]]
+        input_sympy_shape = ctx.get_sympy_shape(node, 0)
+        if get_opset(ctx.out_mp_) <= 10:
+            scales = ctx.try_get_value(node, 1)
+            if scales is not None:
+                new_sympy_shape = [sympy.simplify(sympy.floor(d * s)) for d, s in zip(input_sympy_shape, scales)]
+                ctx.update_computed_dims(new_sympy_shape)
+                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(new_sympy_shape),
+                    )
+                )
+        else:
+            roi = ctx.try_get_value(node, 1)
+            scales = ctx.try_get_value(node, 2)
+            sizes = ctx.try_get_value(node, 3)
+            if sizes is not None:
+                new_sympy_shape = [sympy.simplify(round(s)) for s in sizes]
+                ctx.update_computed_dims(new_sympy_shape)
+            elif scales is not None:
+                rank = len(scales)
+                if get_attribute(node, "coordinate_transformation_mode") == "tf_crop_and_resize":
+                    assert len(roi) == 2 * rank
+                    roi_start = list(roi)[:rank]
+                    roi_end = list(roi)[rank:]
+                else:
+                    roi_start = [0] * rank
+                    roi_end = [1] * rank
+                if isinstance(scales, np.ndarray):
+                    scales = scales.tolist()
+                else:
+                    scales = list(scales)
+                new_sympy_shape = [
+                    sympy.floor(d * (end - start) * scale + sympy.Rational(1, 2))
+                    for d, start, end, scale in zip(input_sympy_shape, roi_start, roi_end, scales)
+                ]
+                ctx.update_computed_dims(new_sympy_shape)
+            else:
+                new_sympy_shape = ctx.new_symbolic_shape(ctx.get_shape_rank(node, 0), node)
+
+            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(new_sympy_shape),
+                )
+            )
+
+
+register_shape_handler(ResizeHandler())

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

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ScatterElements operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class ScatterElementsHandler(ShapeHandler):
+    """Handler for ScatterElements operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ScatterElements"
+
+    def infer_shape(self, node, ctx) -> None:
+        data_shape = ctx.get_shape(node, 0)
+        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,
+            )
+        )
+
+
+register_shape_handler(ScatterElementsHandler())

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

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SoftmaxCrossEntropyLoss operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import MultiOpHandler, ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class SoftmaxCrossEntropyLossHandler(ShapeHandler):
+    """Handler for SoftmaxCrossEntropyLoss operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SoftmaxCrossEntropyLoss"
+
+    def infer_shape(self, node, ctx) -> None:
+        vi = ctx.known_vi_[node.output[0]]
+        elem_type = ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+
+        specified_output_type = get_attribute(node, "output_type", None)
+        if specified_output_type is not None:
+            elem_type = specified_output_type
+
+        vi.type.tensor_type.elem_type = elem_type
+        vi.type.tensor_type.shape.CopyFrom(onnx.TensorShapeProto())
+
+        if len(node.output) > 1:
+            data_shape = ctx.get_shape(node, 0)
+            vi = ctx.known_vi_[node.output[1]]
+            vi.CopyFrom(helper.make_tensor_value_info(vi.name, elem_type, data_shape))
+
+
+def _infer_softmax_cross_entropy(node, ctx):
+    SoftmaxCrossEntropyLossHandler().infer_shape(node, ctx)
+
+
+register_shape_handler(SoftmaxCrossEntropyLossHandler())
+register_shape_handler(MultiOpHandler("SoftmaxCrossEntropyLossInternal", _infer_softmax_cross_entropy))
+register_shape_handler(MultiOpHandler("NegativeLogLikelihoodLossInternal", _infer_softmax_cross_entropy))

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

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for TopK operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import as_scalar, get_attribute, get_opset, get_shape_from_sympy_shape, handle_negative_axis
+
+
+class TopKHandler(ShapeHandler):
+    """Handler for TopK operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "TopK"
+
+    def infer_shape(self, node, ctx) -> None:
+        rank = ctx.get_shape_rank(node, 0)
+        axis = handle_negative_axis(get_attribute(node, "axis", -1), rank)
+        new_shape = ctx.get_shape(node, 0)
+
+        if get_opset(ctx.out_mp_) <= 9:
+            k = get_attribute(node, "k")
+        else:
+            k = ctx.get_int_or_float_values(node)[1]
+
+        k = ctx.new_symbolic_dim_from_output(node) if k is None else as_scalar(k)
+        if type(k) in {int, str}:
+            new_shape[axis] = k
+        else:
+            new_sympy_shape = ctx.get_sympy_shape(node, 0)
+            new_sympy_shape[axis] = k
+            ctx.update_computed_dims(new_sympy_shape)
+            new_shape = get_shape_from_sympy_shape(new_sympy_shape)
+
+        for i_o in range(len(node.output)):
+            vi = ctx.known_vi_[node.output[i_o]]
+            vi.CopyFrom(helper.make_tensor_value_info(node.output[i_o], vi.type.tensor_type.elem_type, new_shape))
+
+
+register_shape_handler(TopKHandler())

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

@@ -0,0 +1,18 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Neural network operator shape handlers."""
+
+from . import conv
+from . import nhwc_conv
+from . import average_pool
+from . import max_pool
+from . import batch_normalization
+from . import identity
+from . import cum_sum
+from . import round
+from . import reciprocal
+from . import memcpy_from_host
+from . import memcpy_to_host
+from . import moe
+from . import all_reduce

onnxslim/core/shape_inference/standard_ops/nn/all_reduce.py

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

onnxslim/core/shape_inference/standard_ops/nn/average_pool.py

@@ -0,0 +1,40 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for AveragePool operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_shape_from_sympy_shape
+
+
+class PoolHandler(ShapeHandler):
+    """Handler for pooling operators (AveragePool, MaxPool)."""
+
+    def __init__(self, op_type_name):
+        super().__init__()
+        self._op_type = op_type_name
+
+    @property
+    def op_type(self) -> str:
+        return self._op_type
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.compute_conv_pool_shape(node)
+        ctx.update_computed_dims(sympy_shape)
+        for o in node.output:
+            if not o:
+                continue
+            vi = ctx.known_vi_[o]
+            vi.CopyFrom(
+                helper.make_tensor_value_info(
+                    o,
+                    vi.type.tensor_type.elem_type,
+                    get_shape_from_sympy_shape(sympy_shape),
+                )
+            )
+
+
+register_shape_handler(PoolHandler("AveragePool"))

onnxslim/core/shape_inference/standard_ops/nn/batch_normalization.py

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for BatchNormalization operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class BatchNormalizationHandler(ShapeHandler):
+    """Handler for BatchNormalization operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "BatchNormalization"
+
+    def infer_shape(self, node, ctx) -> None:
+        ctx.propagate_shape_and_type(node)
+
+        # this works for opsets < 14 and 14 since we check i < len(node.output) in the loop
+        for i in {1, 2, 3, 4}:
+            if i < len(node.output) and node.output[i]:
+                ctx.propagate_shape_and_type(node, input_index=1, output_index=i)
+
+
+register_shape_handler(BatchNormalizationHandler())

onnxslim/core/shape_inference/standard_ops/nn/conv.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Conv operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_shape_from_sympy_shape
+
+
+class ConvHandler(ShapeHandler):
+    """Handler for Conv operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Conv"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.compute_conv_pool_shape(node)
+        ctx.update_computed_dims(sympy_shape)
+        vi = ctx.known_vi_[node.output[0]]
+        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(ConvHandler())

onnxslim/core/shape_inference/standard_ops/nn/cum_sum.py

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

onnxslim/core/shape_inference/standard_ops/nn/identity.py

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

onnxslim/core/shape_inference/standard_ops/nn/max_pool.py

@@ -0,0 +1,9 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for MaxPool operator."""
+
+from ...registry import register_shape_handler
+from .average_pool import PoolHandler
+
+register_shape_handler(PoolHandler("MaxPool"))

onnxslim/core/shape_inference/standard_ops/nn/memcpy_from_host.py

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

onnxslim/core/shape_inference/standard_ops/nn/memcpy_to_host.py

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

onnxslim/core/shape_inference/standard_ops/nn/moe.py

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

onnxslim/core/shape_inference/standard_ops/nn/nhwc_conv.py

@@ -0,0 +1,33 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for NhwcConv operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_shape_from_sympy_shape
+
+
+class NhwcConvHandler(ShapeHandler):
+    """Handler for NhwcConv operator (channels last format)."""
+
+    @property
+    def op_type(self) -> str:
+        return "NhwcConv"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.compute_conv_pool_shape(node, channels_last=True)
+        ctx.update_computed_dims(sympy_shape)
+        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,
+                get_shape_from_sympy_shape(sympy_shape),
+            )
+        )
+
+
+register_shape_handler(NhwcConvHandler())

onnxslim/core/shape_inference/standard_ops/nn/reciprocal.py

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

onnxslim/core/shape_inference/standard_ops/nn/round.py

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

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

@@ -0,0 +1,10 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Sequence operator shape handlers."""
+
+from . import concat_from_sequence
+from . import split_to_sequence
+from . import sequence_at
+from . import sequence_insert
+from . import zip_map

onnxslim/core/shape_inference/standard_ops/sequence/concat_from_sequence.py

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