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

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

@@ -0,0 +1,31 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SequenceAt operator."""
+
+import onnx
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class SequenceAtHandler(ShapeHandler):
+    """Handler for SequenceAt operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SequenceAt"
+
+    def infer_shape(self, node, ctx) -> None:
+        seq_shape = ctx.get_shape(node, 0)
+        if seq_shape is not None:
+            vi = ctx.known_vi_[node.output[0]]
+            for di, d in enumerate(seq_shape):
+                if d is not None:
+                    continue
+                new_dim = onnx.TensorShapeProto.Dimension()
+                new_dim.dim_param = str(ctx.new_symbolic_dim_from_output(node, 0, di))
+                vi.type.tensor_type.shape.dim[di].CopyFrom(new_dim)
+
+
+register_shape_handler(SequenceAtHandler())

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

@@ -0,0 +1,26 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SequenceInsert operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+
+
+class SequenceInsertHandler(ShapeHandler):
+    """Handler for SequenceInsert operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SequenceInsert"
+
+    def infer_shape(self, node, ctx) -> None:
+        vi_seq = ctx.known_vi_[node.input[0]]
+        vi_tensor = ctx.known_vi_[node.input[1]]
+        vi_out_seq = ctx.known_vi_[node.output[0]]
+        vi_out_seq.CopyFrom(vi_seq)
+        vi_out_seq.name = node.output[0]
+        ctx.fuse_tensor_type(node, 0, vi_out_seq.type, vi_tensor.type)
+
+
+register_shape_handler(SequenceInsertHandler())

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

@@ -0,0 +1,24 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for SplitToSequence operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ..tensor.split import infer_split_common
+
+
+class SplitToSequenceHandler(ShapeHandler):
+    """Handler for SplitToSequence operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "SplitToSequence"
+
+    def infer_shape(self, node, ctx) -> None:
+        infer_split_common(node, ctx, helper.make_sequence_value_info)
+
+
+register_shape_handler(SplitToSequenceHandler())

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

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ZipMap operator."""
+
+import onnx
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class ZipMapHandler(ShapeHandler):
+    """Handler for ZipMap operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "ZipMap"
+
+    def infer_shape(self, node, ctx) -> None:
+        map_key_type = None
+        if get_attribute(node, "classlabels_int64s") is not None:
+            map_key_type = onnx.TensorProto.INT64
+        elif get_attribute(node, "classlabels_strings") is not None:
+            map_key_type = onnx.TensorProto.STRING
+
+        assert map_key_type is not None
+        new_vi = onnx.ValueInfoProto()
+        new_vi.name = node.output[0]
+        new_vi.type.sequence_type.elem_type.map_type.value_type.tensor_type.elem_type = onnx.TensorProto.FLOAT
+        new_vi.type.sequence_type.elem_type.map_type.key_type = map_key_type
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(new_vi)
+
+
+register_shape_handler(ZipMapHandler())

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

@@ -0,0 +1,20 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Tensor manipulation operator shape handlers."""
+
+from . import concat
+from . import gather
+from . import gather_elements
+from . import gather_nd
+from . import reshape
+from . import slice
+from . import split
+from . import squeeze
+from . import unsqueeze
+from . import transpose
+from . import tile
+from . import expand
+from . import pad
+from . import shape
+from . import size

onnxslim/core/shape_inference/standard_ops/tensor/concat.py

@@ -0,0 +1,62 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Concat 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
+
+
+class ConcatHandler(ShapeHandler):
+    """Handler for Concat operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Concat"
+
+    def infer_shape(self, node, ctx) -> None:
+        if any(i in ctx.sympy_data_ or i in ctx.initializers_ for i in node.input):
+            values = ctx.get_int_or_float_values(node)
+            if all(v is not None for v in values):
+                assert get_attribute(node, "axis") == 0
+                ctx.sympy_data_[node.output[0]] = []
+                for i in range(len(node.input)):
+                    value = values[i]
+                    if isinstance(value, list):
+                        ctx.sympy_data_[node.output[0]].extend(value)
+                    else:
+                        ctx.sympy_data_[node.output[0]].append(value)
+
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        axis = handle_negative_axis(get_attribute(node, "axis"), len(sympy_shape))
+        for i_idx in range(1, len(node.input)):
+            input_shape = ctx.get_sympy_shape(node, i_idx)
+            if input_shape:
+                sympy_shape[axis] = sympy_shape[axis] + input_shape[axis]
+        ctx.update_computed_dims(sympy_shape)
+        # merge symbolic dims for non-concat axes
+        for d in range(len(sympy_shape)):
+            if d == axis:
+                continue
+            dims = [ctx.get_shape(node, i_idx)[d] for i_idx in range(len(node.input)) if ctx.get_shape(node, i_idx)]
+            if all(d == dims[0] for d in dims):
+                continue
+            merged = ctx.merge_symbols(dims)
+            if type(merged) == str:
+                sympy_shape[d] = ctx.symbolic_dims_[merged] if merged else None
+            else:
+                sympy_shape[d] = merged
+        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(ConcatHandler())

onnxslim/core/shape_inference/standard_ops/tensor/expand.py

@@ -0,0 +1,36 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Expand operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import as_list, get_shape_from_sympy_shape
+
+
+class ExpandHandler(ShapeHandler):
+    """Handler for Expand operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Expand"
+
+    def infer_shape(self, node, ctx) -> None:
+        expand_to_shape = as_list(ctx.try_get_value(node, 1), keep_none=True)
+        if expand_to_shape is not None:
+            ctx.update_computed_dims(expand_to_shape)
+            shape = ctx.get_shape(node, 0)
+            new_shape = ctx.broadcast_shapes(shape, get_shape_from_sympy_shape(expand_to_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,
+                    new_shape,
+                )
+            )
+
+
+register_shape_handler(ExpandHandler())

onnxslim/core/shape_inference/standard_ops/tensor/gather.py

@@ -0,0 +1,48 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Gather operator."""
+
+import numpy as np
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, handle_negative_axis
+
+
+class GatherHandler(ShapeHandler):
+    """Handler for Gather operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Gather"
+
+    def infer_shape(self, node, ctx) -> None:
+        data_shape = ctx.get_shape(node, 0)
+        axis = handle_negative_axis(get_attribute(node, "axis", 0), len(data_shape))
+        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[:axis] + indices_shape + data_shape[axis + 1 :],
+            )
+        )
+        # for 1D input, do some sympy compute
+        if node.input[0] in ctx.sympy_data_ and len(data_shape) == 1 and get_attribute(node, "axis", 0) == 0:
+            idx = ctx.try_get_value(node, 1)
+            if idx is not None:
+                data = ctx.sympy_data_[node.input[0]]
+                if type(data) == list:
+                    if type(idx) == np.ndarray and len(idx.shape) == 1:
+                        ctx.sympy_data_[node.output[0]] = [data[int(i)] for i in idx]
+                    else:
+                        ctx.sympy_data_[node.output[0]] = data[int(idx)]
+                else:
+                    assert idx in {0, -1}
+                    ctx.sympy_data_[node.output[0]] = data
+
+
+register_shape_handler(GatherHandler())

onnxslim/core/shape_inference/standard_ops/tensor/gather_elements.py

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

onnxslim/core/shape_inference/standard_ops/tensor/gather_nd.py

@@ -0,0 +1,42 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for GatherND operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, is_literal
+
+
+class GatherNDHandler(ShapeHandler):
+    """Handler for GatherND operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "GatherND"
+
+    def infer_shape(self, node, ctx) -> None:
+        data_shape = ctx.get_shape(node, 0)
+        data_rank = len(data_shape)
+        indices_shape = ctx.get_shape(node, 1)
+        last_index_dimension = indices_shape[-1]
+        batch_dims = get_attribute(node, "batch_dims", 0)
+        assert (
+            is_literal(last_index_dimension)
+            and is_literal(batch_dims)
+            and (batch_dims + last_index_dimension) <= data_rank
+        )
+        new_shape = indices_shape[:-1] + data_shape[batch_dims + last_index_dimension :]
+        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,
+                new_shape,
+            )
+        )
+
+
+register_shape_handler(GatherNDHandler())

onnxslim/core/shape_inference/standard_ops/tensor/pad.py

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Pad 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
+
+
+class PadHandler(ShapeHandler):
+    """Handler for Pad operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Pad"
+
+    def infer_shape(self, node, ctx) -> None:
+        if get_opset(ctx.out_mp_) <= 10:
+            pads = get_attribute(node, "pads")
+        else:
+            pads = ctx.try_get_value(node, 1)
+
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        rank = len(sympy_shape)
+
+        if pads is not None:
+            assert len(pads) == 2 * rank
+            new_sympy_shape = [d + pad_up + pad_down for d, pad_up, pad_down in zip(sympy_shape, pads[:rank], pads[rank:])]
+            ctx.update_computed_dims(new_sympy_shape)
+        else:
+            new_sympy_shape = ctx.new_symbolic_shape(rank, node)
+        output_tp = 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_tp, get_shape_from_sympy_shape(new_sympy_shape)))
+
+
+register_shape_handler(PadHandler())

onnxslim/core/shape_inference/standard_ops/tensor/reshape.py

@@ -0,0 +1,72 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Reshape operator."""
+
+import sympy
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_shape_from_sympy_shape, is_literal
+
+
+class ReshapeHandler(ShapeHandler):
+    """Handler for Reshape operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Reshape"
+
+    def infer_shape(self, node, ctx) -> None:
+        shape_value = ctx.try_get_value(node, 1)
+        vi = ctx.known_vi_[node.output[0]]
+        if shape_value is None:
+            shape_shape = ctx.get_shape(node, 1)
+            assert len(shape_shape) == 1
+            shape_rank = shape_shape[0]
+            assert is_literal(shape_rank)
+            vi.CopyFrom(
+                helper.make_tensor_value_info(
+                    node.output[0],
+                    vi.type.tensor_type.elem_type,
+                    get_shape_from_sympy_shape(ctx.new_symbolic_shape(shape_rank, node)),
+                )
+            )
+        else:
+            input_sympy_shape = ctx.get_sympy_shape(node, 0)
+            total = 1
+            for d in input_sympy_shape:
+                total = total * d
+            new_sympy_shape = []
+            deferred_dim_idx = -1
+            non_deferred_size = 1
+            for i, d in enumerate(shape_value):
+                if type(d) == sympy.Symbol or d != 0:
+                    new_sympy_shape.append(d)
+                else:
+                    new_sympy_shape.append(input_sympy_shape[i])
+                    non_deferred_size = non_deferred_size * input_sympy_shape[i]
+                if d == -1:
+                    deferred_dim_idx = i
+                elif d != 0:
+                    non_deferred_size = non_deferred_size * d
+
+            assert new_sympy_shape.count(-1) < 2
+            if -1 in new_sympy_shape:
+                new_dim = total // non_deferred_size
+                new_sympy_shape[deferred_dim_idx] = new_dim
+
+            ctx.update_computed_dims(new_sympy_shape)
+            vi.CopyFrom(
+                helper.make_tensor_value_info(
+                    node.output[0],
+                    vi.type.tensor_type.elem_type,
+                    get_shape_from_sympy_shape(new_sympy_shape),
+                )
+            )
+
+        ctx.pass_on_sympy_data(node)
+
+
+register_shape_handler(ReshapeHandler())

onnxslim/core/shape_inference/standard_ops/tensor/shape.py

@@ -0,0 +1,38 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Shape operator."""
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class ShapeOpHandler(ShapeHandler):
+    """Handler for Shape operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Shape"
+
+    def infer_shape(self, node, ctx) -> None:
+        start = get_attribute(node, "start", 0)
+        end = get_attribute(node, "end", None)
+
+        full_sympy_shape = ctx.get_sympy_shape(node, 0)
+        num_dims = len(full_sympy_shape)
+
+        if start < 0:
+            start = num_dims + start
+        if end is None:
+            end = num_dims
+        elif end < 0:
+            end = num_dims + end
+
+        assert 0 <= start <= end <= num_dims, f"reshape start/end invalid: start={start}, end={end}, total_dims={num_dims}"
+
+        target_sympy_shape = full_sympy_shape[start:end]
+        ctx.sympy_data_[node.output[0]] = target_sympy_shape
+
+
+register_shape_handler(ShapeOpHandler())

onnxslim/core/shape_inference/standard_ops/tensor/size.py

@@ -0,0 +1,29 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Size operator."""
+
+import onnx
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import sympy_reduce_product
+
+
+class SizeHandler(ShapeHandler):
+    """Handler for Size operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Size"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        ctx.sympy_data_[node.output[0]] = sympy_reduce_product(sympy_shape)
+        ctx.known_vi_[node.output[0]].CopyFrom(
+            helper.make_tensor_value_info(node.output[0], onnx.TensorProto.INT64, [])
+        )
+
+
+register_shape_handler(SizeHandler())