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

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

@@ -0,0 +1,183 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Slice operator."""
+
+import logging
+
+import numpy as np
+import sympy
+from onnx import helper
+
+from onnxslim.third_party._sympy.solve import try_solve
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import (
+    as_list,
+    get_attribute,
+    get_opset,
+    get_shape_from_sympy_shape,
+    is_literal,
+)
+
+logger = logging.getLogger(__name__)
+
+
+class SliceHandler(ShapeHandler):
+    """Handler for Slice operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Slice"
+
+    def infer_shape(self, node, ctx) -> None:
+        def flatten_min(expr):
+            """Returns a list with expressions split by min() for inequality proof."""
+            assert isinstance(expr, sympy.Add), f"Expected a sum of two arguments, got {expr}"
+            min_positions = [idx for idx in range(len(expr.args)) if isinstance(expr.args[idx], sympy.Min)]
+            if len(min_positions) == 1:
+                min_pos = min_positions[0]
+
+                def replace_min_with_arg(arg_idx):
+                    replaced = list(expr.args)
+                    assert isinstance(replaced[min_pos], sympy.Min)
+                    assert len(replaced[min_pos].args) == 2
+                    replaced[min_pos] = replaced[min_pos].args[arg_idx]
+                    return sympy.Add(*replaced)
+
+                return [replace_min_with_arg(0), replace_min_with_arg(1)]
+            return [expr]
+
+        def less_equal(x, y):
+            """Returns True if x is less than or equal to y."""
+            try:
+                return x <= y
+            except TypeError:
+                pass
+            try:
+                return y >= x
+            except TypeError:
+                pass
+            try:
+                return -x >= -y
+            except TypeError:
+                pass
+            try:
+                return -y <= -x
+            except TypeError:
+                pass
+            try:
+                return y - x >= 0
+            except TypeError:
+                return all(d >= 0 for d in flatten_min(y - x))
+
+        def handle_negative_index(index, bound):
+            """Normalizes a negative index to be in [0, bound)."""
+            try:
+                if not less_equal(0, index):
+                    if is_literal(index) and index <= -ctx.int_max_:
+                        return index
+                    return bound + index
+            except TypeError:
+                logger.warning(f"Cannot determine if {index} < 0")
+            return index
+
+        if get_opset(ctx.out_mp_) <= 9:
+            axes = get_attribute(node, "axes")
+            starts = get_attribute(node, "starts")
+            ends = get_attribute(node, "ends")
+            if not axes:
+                axes = list(range(len(starts)))
+            steps = [1] * len(axes)
+        else:
+            starts = as_list(ctx.try_get_value(node, 1), keep_none=True)
+            ends = as_list(ctx.try_get_value(node, 2), keep_none=True)
+            axes = ctx.try_get_value(node, 3)
+            steps = ctx.try_get_value(node, 4)
+            if axes is None and (starts is not None or ends is not None):
+                axes = list(range(len(starts if starts is not None else ends)))
+            if steps is None and (starts is not None or ends is not None):
+                steps = [1] * len(starts if starts is not None else ends)
+            axes = as_list(axes, keep_none=True)
+            steps = as_list(steps, keep_none=True)
+
+        new_sympy_shape = ctx.get_sympy_shape(node, 0)
+        if starts is None or ends is None:
+            if axes is None:
+                for i in range(len(new_sympy_shape)):
+                    new_sympy_shape[i] = ctx.new_symbolic_dim_from_output(node, 0, i)
+            else:
+                new_sympy_shape = get_shape_from_sympy_shape(new_sympy_shape)
+                for i in axes:
+                    new_sympy_shape[i] = ctx.new_symbolic_dim_from_output(node, 0, i)
+        else:
+            for i, s, e, t in zip(axes, starts, ends, steps):
+                if is_literal(e):
+                    e = handle_negative_index(e, new_sympy_shape[i])
+                if is_literal(e):
+                    if e >= ctx.int_max_:
+                        e = new_sympy_shape[i]
+                    elif e <= -ctx.int_max_:
+                        e = 0 if s > 0 else -1
+                    elif is_literal(new_sympy_shape[i]):
+                        if e < 0:
+                            e = max(0, e + new_sympy_shape[i])
+                        e = min(e, new_sympy_shape[i])
+                    else:
+                        if e > 0:
+                            e = sympy.Min(e, new_sympy_shape[i]) if e > 1 else e
+                else:
+                    if is_literal(new_sympy_shape[i]):
+                        if new_sympy_shape[i] < 0:
+                            e = sympy.Min(e, new_sympy_shape[i])
+                    else:
+                        try:
+                            if not less_equal(e, new_sympy_shape[i]):
+                                e = new_sympy_shape[i]
+                        except Exception:
+                            if len(e.free_symbols) == 1:
+                                if try_solve((e - new_sympy_shape[i]) >= 0, next(iter(e.free_symbols))) is None:
+                                    logger.warning(
+                                        f"Unable to solve if {e} <= {new_sympy_shape[i]}, treat as not equal"
+                                    )
+                            else:
+                                logger.warning(f"Unable to determine if {e} <= {new_sympy_shape[i]}, treat as equal")
+                                e = new_sympy_shape[i]
+
+                s = handle_negative_index(s, new_sympy_shape[i])
+                if is_literal(new_sympy_shape[i]) and is_literal(s):
+                    s = max(0, min(s, new_sympy_shape[i]))
+
+                new_sympy_shape[i] = sympy.simplify((e - s + t + (-1 if t > 0 else 1)) // t)
+
+            ctx.update_computed_dims(new_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(new_sympy_shape),
+            )
+        )
+
+        # handle sympy_data if needed, for slice in shape computation
+        if (
+            node.input[0] in ctx.sympy_data_
+            and [0] == axes
+            and starts is not None
+            and len(starts) == 1
+            and ends is not None
+            and len(ends) == 1
+            and steps is not None
+            and len(steps) == 1
+        ):
+            input_sympy_data = ctx.sympy_data_[node.input[0]]
+            if type(input_sympy_data) == list or (
+                type(input_sympy_data) == np.array and len(input_sympy_data.shape) == 1
+            ):
+                ctx.sympy_data_[node.output[0]] = input_sympy_data[starts[0] : ends[0] : steps[0]]
+
+
+register_shape_handler(SliceHandler())

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

@@ -0,0 +1,57 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Split operator."""
+
+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, handle_negative_axis
+
+
+class SplitHandler(ShapeHandler):
+    """Handler for Split operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Split"
+
+    def infer_shape(self, node, ctx) -> None:
+        infer_split_common(node, ctx, helper.make_tensor_value_info)
+
+
+def infer_split_common(node, ctx, make_value_info_func):
+    """Infers the output shape for the Split operator."""
+    input_sympy_shape = ctx.get_sympy_shape(node, 0)
+    axis = handle_negative_axis(get_attribute(node, "axis", 0), len(input_sympy_shape))
+    op_set = get_opset(ctx.out_mp_)
+
+    if op_set < 13:
+        split = get_attribute(node, "split")
+        assert ctx.try_get_value(node, 1) is None
+    else:
+        split = ctx.try_get_value(node, 1)
+        assert get_attribute(node, "split") is None
+
+    if split is None:
+        num_outputs = len(node.output)
+        split = [input_sympy_shape[axis] / sympy.Integer(num_outputs)] * num_outputs
+        ctx.update_computed_dims(split)
+    else:
+        split = [sympy.Integer(s) for s in split]
+
+    for i_o in range(len(split)):
+        vi = ctx.known_vi_[node.output[i_o]]
+        vi.CopyFrom(
+            make_value_info_func(
+                node.output[i_o],
+                ctx.known_vi_[node.input[0]].type.tensor_type.elem_type,
+                get_shape_from_sympy_shape([*input_sympy_shape[:axis], split[i_o], *input_sympy_shape[axis + 1 :]]),
+            )
+        )
+        ctx.known_vi_[vi.name] = vi
+
+
+register_shape_handler(SplitHandler())

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

@@ -0,0 +1,69 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Squeeze operator."""
+
+import logging
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_opset, handle_negative_axis
+
+logger = logging.getLogger(__name__)
+
+
+class SqueezeHandler(ShapeHandler):
+    """Handler for Squeeze operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Squeeze"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_shape = ctx.get_shape(node, 0)
+        op_set = get_opset(ctx.out_mp_)
+
+        if op_set < 13:
+            axes = get_attribute(node, "axes")
+            assert ctx.try_get_value(node, 1) is None
+        else:
+            axes = ctx.try_get_value(node, 1)
+            assert get_attribute(node, "axes") is None
+
+        if axes is None:
+            output_shape = [s for s in input_shape if s != 1]
+            if ctx.verbose_ > 0:
+                symbolic_dimensions = [s for s in input_shape if type(s) != int]
+                if symbolic_dimensions:
+                    logger.debug(
+                        f"Symbolic dimensions in input shape of op: '{node.op_type}' node: '{node.name}'. "
+                        f"Assuming the following dimensions are never equal to 1: {symbolic_dimensions}"
+                    )
+        else:
+            axes = [handle_negative_axis(a, len(input_shape)) for a in axes]
+            output_shape = []
+            for i in range(len(input_shape)):
+                if i not in axes:
+                    output_shape.append(input_shape[i])
+                else:
+                    assert input_shape[i] == 1 or type(input_shape[i]) != int
+                    if ctx.verbose_ > 0 and type(input_shape[i]) != int:
+                        logger.debug(
+                            f"Symbolic dimensions in input shape of op: '{node.op_type}' node: '{node.name}'. "
+                            f"Assuming the dimension '{input_shape[i]}' at index {i} of the input to be equal to 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,
+                output_shape,
+            )
+        )
+        ctx.pass_on_sympy_data(node)
+
+
+register_shape_handler(SqueezeHandler())

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

@@ -0,0 +1,41 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Tile operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_shape_from_sympy_shape
+
+
+class TileHandler(ShapeHandler):
+    """Handler for Tile operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Tile"
+
+    def infer_shape(self, node, ctx) -> None:
+        repeats_value = ctx.try_get_value(node, 1)
+        new_sympy_shape = []
+        if repeats_value is not None:
+            input_sympy_shape = ctx.get_sympy_shape(node, 0)
+            for i, d in enumerate(input_sympy_shape):
+                new_dim = d * repeats_value[i]
+                new_sympy_shape.append(new_dim)
+            ctx.update_computed_dims(new_sympy_shape)
+        else:
+            new_sympy_shape = ctx.new_symbolic_shape(ctx.get_shape_rank(node, 0), node)
+        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(new_sympy_shape),
+            )
+        )
+
+
+register_shape_handler(TileHandler())

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

@@ -0,0 +1,30 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Transpose operator."""
+
+import numpy as np
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute
+
+
+class TransposeHandler(ShapeHandler):
+    """Handler for Transpose operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Transpose"
+
+    def infer_shape(self, node, ctx) -> None:
+        if node.input[0] in ctx.sympy_data_:
+            data_shape = ctx.get_shape(node, 0)
+            perm = get_attribute(node, "perm", reversed(list(range(len(data_shape)))))
+            input_data = ctx.sympy_data_[node.input[0]]
+            ctx.sympy_data_[node.output[0]] = (
+                np.transpose(np.array(input_data).reshape(*data_shape), axes=tuple(perm)).flatten().tolist()
+            )
+
+
+register_shape_handler(TransposeHandler())

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

@@ -0,0 +1,54 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for Unsqueeze operator."""
+
+from onnx import helper
+
+from ...base import ShapeHandler
+from ...registry import register_shape_handler
+from ...utils import get_attribute, get_opset, handle_negative_axis
+
+
+class UnsqueezeHandler(ShapeHandler):
+    """Handler for Unsqueeze operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "Unsqueeze"
+
+    def infer_shape(self, node, ctx) -> None:
+        input_shape = ctx.get_shape(node, 0)
+        op_set = get_opset(ctx.out_mp_)
+
+        if op_set < 13:
+            axes = get_attribute(node, "axes")
+            assert ctx.try_get_value(node, 1) is None
+        else:
+            axes = ctx.try_get_value(node, 1)
+            assert get_attribute(node, "axes") is None
+
+        output_rank = len(input_shape) + len(axes)
+        axes = [handle_negative_axis(a, output_rank) for a in axes]
+
+        input_axis = 0
+        output_shape = []
+        for i in range(output_rank):
+            if i in axes:
+                output_shape.append(1)
+            else:
+                output_shape.append(input_shape[input_axis])
+                input_axis += 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,
+                output_shape,
+            )
+        )
+        ctx.pass_on_sympy_data(node)
+
+
+register_shape_handler(UnsqueezeHandler())

onnxslim/core/shape_inference/utils.py

@@ -0,0 +1,244 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Utility functions for symbolic shape inference."""
+
+import numpy as np
+import onnx
+import sympy
+from onnx import helper, numpy_helper
+
+
+def get_attribute(node, attr_name, default_value=None):
+    """Retrieve the value of an attribute from an ONNX node.
+
+    Args:
+        node: The ONNX node.
+        attr_name: The name of the attribute to retrieve.
+        default_value: The default value if the attribute is not found.
+
+    Returns:
+        The attribute value or the default value.
+    """
+    found = [attr for attr in node.attribute if attr.name == attr_name]
+    return helper.get_attribute_value(found[0]) if found else default_value
+
+
+def get_dim_from_proto(dim):
+    """Retrieve the dimension value from the ONNX protobuf object.
+
+    Args:
+        dim: The ONNX TensorShapeProto.Dimension.
+
+    Returns:
+        The dimension value (int or str) or None.
+    """
+    return getattr(dim, dim.WhichOneof("value")) if type(dim.WhichOneof("value")) is str else None
+
+
+def is_sequence(type_proto):
+    """Check if the given ONNX proto type is a sequence.
+
+    Args:
+        type_proto: The ONNX TypeProto.
+
+    Returns:
+        True if the type is a sequence type.
+    """
+    cls_type = type_proto.WhichOneof("value")
+    assert cls_type in {"tensor_type", "sequence_type"}
+    return cls_type == "sequence_type"
+
+
+def get_shape_from_type_proto(type_proto):
+    """Extract the shape of a tensor from an ONNX type proto.
+
+    Args:
+        type_proto: The ONNX TypeProto.
+
+    Returns:
+        A list of dimension values or None if no shape is available.
+    """
+    assert not is_sequence(type_proto)
+    if type_proto.tensor_type.HasField("shape"):
+        return [get_dim_from_proto(d) for d in type_proto.tensor_type.shape.dim]
+    else:
+        return None
+
+
+def get_elem_type_from_type_proto(type_proto):
+    """Return the element type from a given TypeProto object.
+
+    Args:
+        type_proto: The ONNX TypeProto.
+
+    Returns:
+        The element type (e.g., TensorProto.FLOAT).
+    """
+    if is_sequence(type_proto):
+        return type_proto.sequence_type.elem_type.tensor_type.elem_type
+    else:
+        return type_proto.tensor_type.elem_type
+
+
+def get_shape_from_value_info(vi):
+    """Return the shape from the given ValueInfoProto object.
+
+    Args:
+        vi: The ONNX ValueInfoProto.
+
+    Returns:
+        A list of dimension values or None.
+    """
+    cls_type = vi.type.WhichOneof("value")
+    if cls_type is None:
+        return None
+    if not is_sequence(vi.type):
+        return get_shape_from_type_proto(vi.type)
+    if vi.type.sequence_type.elem_type.WhichOneof("value") == "tensor_type":
+        return get_shape_from_type_proto(vi.type.sequence_type.elem_type)
+    else:
+        return None
+
+
+def make_named_value_info(name):
+    """Create and return an ONNX ValueInfoProto object with the specified name.
+
+    Args:
+        name: The name for the ValueInfoProto.
+
+    Returns:
+        A new ValueInfoProto with the given name.
+    """
+    vi = onnx.ValueInfoProto()
+    vi.name = name
+    return vi
+
+
+def get_shape_from_sympy_shape(sympy_shape):
+    """Convert a sympy shape to a list with int, str, or None elements.
+
+    Args:
+        sympy_shape: A list of sympy expressions.
+
+    Returns:
+        A list of int, str, or None values.
+    """
+    return [None if i is None else (int(i) if is_literal(i) else str(i)) for i in sympy_shape]
+
+
+def is_literal(dim):
+    """Check if a dimension is a literal number.
+
+    Args:
+        dim: The dimension value to check.
+
+    Returns:
+        True if the dimension is a literal number.
+    """
+    return type(dim) in {int, np.int64, np.int32, sympy.Integer} or (hasattr(dim, "is_number") and dim.is_number)
+
+
+def handle_negative_axis(axis, rank):
+    """Convert a potentially negative axis to a positive axis.
+
+    Args:
+        axis: The axis value (can be negative).
+        rank: The total rank of the tensor.
+
+    Returns:
+        A non-negative axis value.
+    """
+    assert axis < rank and axis >= -rank
+    return axis if axis >= 0 else rank + axis
+
+
+def get_opset(mp, domain=None):
+    """Retrieve the opset version for a given model namespace.
+
+    Args:
+        mp: The ONNX ModelProto.
+        domain: The domain(s) to check. Defaults to common ONNX domains.
+
+    Returns:
+        The opset version or None if not found.
+    """
+    domain = domain or ["", "onnx", "ai.onnx"]
+    if type(domain) != list:
+        domain = [domain]
+    for opset in mp.opset_import:
+        if opset.domain in domain:
+            return opset.version
+    return None
+
+
+def as_scalar(x):
+    """Convert input to scalar if input is a list with a single item or a NumPy ndarray.
+
+    Args:
+        x: The input value.
+
+    Returns:
+        A scalar value.
+    """
+    if type(x) == list:
+        assert len(x) == 1
+        return x[0]
+    elif type(x) == np.ndarray:
+        return x.item()
+    else:
+        return x
+
+
+def as_list(x, keep_none):
+    """Convert input to list, optionally preserving None values.
+
+    Args:
+        x: The input value.
+        keep_none: If True, return None as-is instead of wrapping in list.
+
+    Returns:
+        A list or None.
+    """
+    if type(x) == list:
+        return x
+    elif type(x) == np.ndarray:
+        return list(x)
+    elif keep_none and x is None:
+        return None
+    else:
+        return [x]
+
+
+def sympy_reduce_product(x):
+    """Reduce a list or element to a product using Sympy's Integer.
+
+    Args:
+        x: A list or single value.
+
+    Returns:
+        The product as a sympy expression.
+    """
+    if type(x) == list:
+        value = sympy.Integer(1)
+        for v in x:
+            value = value * v
+    else:
+        value = x
+    return value
+
+
+def numpy_to_sympy(array):
+    """Convert a numpy array to a list of sympy values.
+
+    Args:
+        array: A numpy array.
+
+    Returns:
+        The converted list or value.
+    """
+    if isinstance(array, np.ndarray):
+        if array.ndim == 0:
+            return int(array.item())
+        return [int(x) for x in array.flatten()]
+    return array