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

onnxslim/core/shape_inference/aten_ops/multinomial.py

@@ -0,0 +1,39 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ATen multinomial operator."""
+
+import onnx
+from onnx import helper
+
+from ..base import ShapeHandler
+from ..registry import register_aten_handler
+from ..utils import get_shape_from_sympy_shape
+
+
+class AtenMultinomialHandler(ShapeHandler):
+    """Handler for ATen multinomial operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "multinomial"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        rank = len(sympy_shape)
+        assert rank in {1, 2}
+        num_samples = ctx.try_get_value(node, 1)
+        di = rank - 1
+        last_dim = num_samples or str(ctx.new_symbolic_dim_from_output(node, 0, di))
+        output_shape = [*sympy_shape[:-1], last_dim]
+        vi = ctx.known_vi_[node.output[0]]
+        vi.CopyFrom(
+            helper.make_tensor_value_info(
+                node.output[0],
+                onnx.TensorProto.INT64,
+                get_shape_from_sympy_shape(output_shape),
+            )
+        )
+
+
+register_aten_handler(AtenMultinomialHandler())

onnxslim/core/shape_inference/aten_ops/numpy_t.py

@@ -0,0 +1,22 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ATen numpy_T operator."""
+
+from ..base import ShapeHandler
+from ..registry import register_aten_handler
+from ..standard_ops.tensor.transpose import TransposeHandler
+
+
+class AtenNumpyTHandler(ShapeHandler):
+    """Handler for ATen numpy_T operator (reuses Transpose logic)."""
+
+    @property
+    def op_type(self) -> str:
+        return "numpy_T"
+
+    def infer_shape(self, node, ctx) -> None:
+        TransposeHandler().infer_shape(node, ctx)
+
+
+register_aten_handler(AtenNumpyTHandler())

onnxslim/core/shape_inference/aten_ops/pool2d.py

@@ -0,0 +1,40 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ATen pooling operators."""
+
+import onnx
+from onnx import helper
+
+from ..base import ShapeHandler
+from ..registry import register_aten_handler
+from ..utils import get_shape_from_sympy_shape
+
+
+class AtenPool2dHandler(ShapeHandler):
+    """Handler for ATen pooling operators (max_pool2d_with_indices, avg_pool2d, _adaptive_avg_pool2d)."""
+
+    def __init__(self, op_name):
+        super().__init__()
+        self._op_type = op_name
+
+    @property
+    def op_type(self) -> str:
+        return self._op_type
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        assert len(sympy_shape) == 4
+        sympy_shape[-2:] = [ctx.new_symbolic_dim_from_output(node, 0, i) for i in {2, 3}]
+        ctx.update_computed_dims(sympy_shape)
+        for i, o in enumerate(node.output):
+            if not o:
+                continue
+            vi = ctx.known_vi_[o]
+            elem_type = onnx.TensorProto.INT64 if i == 1 else ctx.known_vi_[node.input[0]].type.tensor_type.elem_type
+            vi.CopyFrom(helper.make_tensor_value_info(o, elem_type, get_shape_from_sympy_shape(sympy_shape)))
+
+
+register_aten_handler(AtenPool2dHandler("max_pool2d_with_indices"))
+register_aten_handler(AtenPool2dHandler("avg_pool2d"))
+register_aten_handler(AtenPool2dHandler("_adaptive_avg_pool2d"))

onnxslim/core/shape_inference/aten_ops/unfold.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ATen unfold operator."""
+
+from onnx import helper
+
+from ..base import ShapeHandler
+from ..registry import register_aten_handler
+from ..utils import get_shape_from_sympy_shape
+
+
+class AtenUnfoldHandler(ShapeHandler):
+    """Handler for ATen unfold operator."""
+
+    @property
+    def op_type(self) -> str:
+        return "unfold"
+
+    def infer_shape(self, node, ctx) -> None:
+        sympy_shape = ctx.get_sympy_shape(node, 0)
+        dimension = ctx.try_get_value(node, 1)
+        size = ctx.try_get_value(node, 2)
+        step = ctx.try_get_value(node, 3)
+        if dimension is not None and size is not None and step is not None:
+            assert dimension < len(sympy_shape)
+            sympy_shape[dimension] = (sympy_shape[dimension] - size) // step + 1
+            sympy_shape.append(size)
+        else:
+            rank = len(sympy_shape)
+            sympy_shape = ctx.new_symbolic_shape(rank + 1, node)
+        ctx.update_computed_dims(sympy_shape)
+        if node.output[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,
+                    get_shape_from_sympy_shape(sympy_shape),
+                )
+            )
+
+
+register_aten_handler(AtenUnfoldHandler())

onnxslim/core/shape_inference/aten_ops/upsample.py

@@ -0,0 +1,44 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Shape handler for ATen upsample operators."""
+
+import numpy as np
+from onnx import helper
+
+from ..base import ShapeHandler
+from ..registry import register_aten_handler
+
+
+class AtenUpsampleHandler(ShapeHandler):
+    """Handler for ATen upsample operators."""
+
+    def __init__(self, op_name):
+        super().__init__()
+        self._op_type = op_name
+
+    @property
+    def op_type(self) -> str:
+        return self._op_type
+
+    def infer_shape(self, node, ctx) -> None:
+        new_shape = None
+        input_shape = ctx.get_shape(node, 0)
+        if input_shape is not None:
+            new_shape = input_shape[:2]
+            output_size = ctx.try_get_value(node, 1)
+            if output_size is not None:
+                new_shape += [dim_size.item() if type(dim_size) == np.int64 else dim_size for dim_size in output_size]
+            else:
+                rank = len(input_shape)
+                new_shape += [str(ctx.new_symbolic_dim_from_output(node, 0, i)) for i in range(2, rank)]
+        if node.output[0] and new_shape is not None:
+            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_aten_handler(AtenUpsampleHandler("upsample_nearest1d"))
+register_aten_handler(AtenUpsampleHandler("upsample_nearest2d"))
+register_aten_handler(AtenUpsampleHandler("upsample_nearest3d"))
+register_aten_handler(AtenUpsampleHandler("upsample_bicubic2d"))

onnxslim/core/shape_inference/base.py

@@ -0,0 +1,111 @@
+# Copyright (c) Microsoft Corporation. All rights reserved.
+# Licensed under the MIT License.
+
+"""Base classes for shape inference handlers."""
+
+from abc import ABC, abstractmethod
+
+
+class ShapeHandler(ABC):
+    """Abstract base class for shape inference handlers.
+
+    Each handler is responsible for inferring the output shapes for a specific
+    ONNX operator type.
+    """
+
+    def __init__(self, min_opset=1, max_opset=999):
+        """Initialize the shape handler.
+
+        Args:
+            min_opset: Minimum ONNX opset version this handler supports.
+            max_opset: Maximum ONNX opset version this handler supports.
+        """
+        self.min_opset = min_opset
+        self.max_opset = max_opset
+
+    @property
+    @abstractmethod
+    def op_type(self) -> str:
+        """Return the ONNX operator type this handler supports.
+
+        Returns:
+            The operator type string (e.g., "Reshape", "MatMul").
+        """
+        raise NotImplementedError
+
+    def supports_opset(self, opset: int) -> bool:
+        """Check if this handler supports the given opset version.
+
+        Args:
+            opset: The ONNX opset version to check.
+
+        Returns:
+            True if the handler supports this opset version.
+        """
+        return self.min_opset <= opset <= self.max_opset
+
+    @abstractmethod
+    def infer_shape(self, node, ctx) -> None:
+        """Infer the output shapes for the given node.
+
+        Args:
+            node: The ONNX node to infer shapes for.
+            ctx: The InferenceContext providing access to shape information.
+        """
+        raise NotImplementedError
+
+
+class PassthroughHandler(ShapeHandler):
+    """Handler for operators that pass through input shape to output unchanged.
+
+    This is used for operators like Identity, Reciprocal, Round, etc.
+    """
+
+    def __init__(self, op_type_name, min_opset=1, max_opset=999):
+        """Initialize the passthrough handler.
+
+        Args:
+            op_type_name: The operator type name.
+            min_opset: Minimum ONNX opset version this handler supports.
+            max_opset: Maximum ONNX opset version this handler supports.
+        """
+        super().__init__(min_opset, max_opset)
+        self._op_type = op_type_name
+
+    @property
+    def op_type(self) -> str:
+        """Return the ONNX operator type this handler supports."""
+        return self._op_type
+
+    def infer_shape(self, node, ctx) -> None:
+        """Pass through shape and type from input to output."""
+        ctx.pass_on_shape_and_type(node)
+
+
+class MultiOpHandler(ShapeHandler):
+    """Handler that supports multiple operator types with the same logic.
+
+    This is useful when multiple operators share the same shape inference logic.
+    """
+
+    def __init__(self, op_type_name, handler_func, min_opset=1, max_opset=999):
+        """Initialize the multi-op handler.
+
+        Args:
+            op_type_name: The operator type name.
+            handler_func: The function to call for shape inference.
+            min_opset: Minimum ONNX opset version this handler supports.
+            max_opset: Maximum ONNX opset version this handler supports.
+        """
+        super().__init__(min_opset, max_opset)
+        self._op_type = op_type_name
+        self._handler_func = handler_func
+
+    @property
+    def op_type(self) -> str:
+        """Return the ONNX operator type this handler supports."""
+        return self._op_type
+
+    def infer_shape(self, node, ctx) -> None:
+        """Call the handler function for shape inference."""
+        self._handler_func(node, ctx)