Add Squeeze and Unsqueeze operator support across quantizer and circuit builder

python/core/model_processing/converters/onnx_converter.py

@@ -1022,7 +1022,9 @@ def _extract_model_io_info(
         ]
         self.input_shape = get_input_shapes(onnx_model)
 
-    def get_weights(self: ONNXConverter) -> tuple[
+    def get_weights(
+        self: ONNXConverter,
+    ) -> tuple[
         dict[str, list[ONNXLayerDict]],
         dict[str, list[ONNXLayerDict]],
         CircuitParamsDict,
@@ -1055,7 +1057,7 @@ def get_weights(self: ONNXConverter) -> tuple[
             scale_base=scale_base,
         )
         # Get layers in correct format
-        (architecture, w_and_b) = self.analyze_layers(
+        architecture, w_and_b = self.analyze_layers(
             scaled_and_transformed_model,
             output_name_to_shape,
         )

python/core/model_processing/onnx_quantizer/layers/squeeze.py

@@ -0,0 +1,155 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, ClassVar
+
+import numpy as np
+from onnx import numpy_helper
+
+from python.core.model_processing.onnx_quantizer.exceptions import (
+    InvalidParamError,
+)
+from python.core.model_processing.onnx_quantizer.layers.base import (
+    BaseOpQuantizer,
+    QuantizerBase,
+    ScaleConfig,
+)
+
+if TYPE_CHECKING:
+    import onnx
+
+
+class QuantizeSqueeze(QuantizerBase):
+    OP_TYPE = "Squeeze"
+    DOMAIN = ""
+    USE_WB = False
+    USE_SCALING = False
+    # Only the data input is relevant for scale-planning.
+    SCALE_PLAN: ClassVar = {0: 1}
+
+
+class SqueezeQuantizer(BaseOpQuantizer, QuantizeSqueeze):
+    """
+    Quantizer for ONNX Squeeze.
+
+    Squeeze is scale-preserving (pure shape/view transform):
+    - No arithmetic
+    - No rescaling
+    - No custom op
+
+    We support:
+    - axes as an attribute (older opsets)
+    - axes as a constant initializer input (newer opsets)
+
+    We do NOT support dynamic axes provided at runtime.
+    """
+
+    def __init__(
+        self: SqueezeQuantizer,
+        new_initializers: list[onnx.TensorProto] | None = None,
+    ) -> None:
+        super().__init__()
+        if new_initializers is not None:
+            self.new_initializers = new_initializers
+
+    def quantize(
+        self: SqueezeQuantizer,
+        node: onnx.NodeProto,
+        graph: onnx.GraphProto,
+        scale_config: ScaleConfig,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[onnx.NodeProto]:
+        # Pure passthrough; QuantizerBase handles standard bookkeeping.
+        return QuantizeSqueeze.quantize(
+            self,
+            node,
+            graph,
+            scale_config,
+            initializer_map,
+        )
+
+    _N_INPUTS_NO_AXES: ClassVar[int] = 1
+    _N_INPUTS_WITH_AXES: ClassVar[int] = 2
+
+    def _get_axes_from_attribute(self, node: onnx.NodeProto) -> list[int] | None:
+        for attr in node.attribute:
+            if attr.name == "axes":
+                return list(attr.ints)
+        return None
+
+    def _get_axes_from_initializer_input(
+        self,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto],
+    ) -> list[int]:
+        axes_name = node.input[1]
+        if axes_name not in initializer_map:
+            raise InvalidParamError(
+                node_name=node.name,
+                op_type=node.op_type,
+                message=(
+                    f"Dynamic axes input is not supported for Squeeze "
+                    f"(expected axes '{axes_name}' to be an initializer)."
+                ),
+            )
+
+        axes_tensor = initializer_map[axes_name]
+        arr = numpy_helper.to_array(axes_tensor)
+
+        if not np.issubdtype(arr.dtype, np.integer):
+            raise InvalidParamError(
+                node_name=node.name,
+                op_type=node.op_type,
+                message=f"Squeeze axes initializer must be integer, got {arr.dtype}.",
+                attr_key="axes",
+                expected="integer tensor (0-D or 1-D)",
+            )
+
+        if arr.ndim == 0:
+            return [int(arr)]
+        if arr.ndim == 1:
+            return [int(x) for x in arr.tolist()]
+
+        raise InvalidParamError(
+            node_name=node.name,
+            op_type=node.op_type,
+            message=f"Squeeze axes initializer must be 0-D or 1-D, got {arr.ndim}-D.",
+            attr_key="axes",
+            expected="0-D scalar or 1-D list of axes",
+        )
+
+    def check_supported(
+        self: SqueezeQuantizer,
+        node: onnx.NodeProto,
+        initializer_map: dict[str, onnx.TensorProto] | None = None,
+    ) -> None:
+        self.validate_node_has_output(node)
+        initializer_map = initializer_map or {}
+
+        n_inputs = len(node.input)
+        if n_inputs not in (self._N_INPUTS_NO_AXES, self._N_INPUTS_WITH_AXES):
+            raise InvalidParamError(
+                node_name=node.name,
+                op_type=node.op_type,
+                message=f"Squeeze expects 1 or 2 inputs, got {n_inputs}.",
+            )
+
+        axes = self._get_axes_from_attribute(node)
+
+        # If axes is provided as a second input, it must be a constant initializer.
+        if axes is None and n_inputs == self._N_INPUTS_WITH_AXES:
+            axes = self._get_axes_from_initializer_input(node, initializer_map)
+
+        # If axes is omitted entirely, ONNX semantics are "remove all dims of size 1".
+        # We can't validate legality here without rank/shape; defer to Rust.
+
+        if axes is None:
+            return
+
+        if len(set(axes)) != len(axes):
+            raise InvalidParamError(
+                node_name=node.name,
+                op_type=node.op_type,
+                message=f"axes must not contain duplicates: {axes}",
+                attr_key="axes",
+                expected="axes list with unique entries",
+            )