Initialisation du repository de Beta

venv/lib/python3.12/site-packages/onnxruntime/quantization/operators/where.py

@@ -0,0 +1,87 @@
+import onnx
+
+from ..quant_utils import TENSOR_NAME_QUANT_SUFFIX, QuantizedValue, QuantizedValueType, attribute_to_kwarg, ms_domain
+from .base_operator import QuantOperatorBase
+from .qdq_base_operator import QDQOperatorBase
+
+
+class QLinearWhere(QuantOperatorBase):
+    def should_quantize(self):
+        return True
+
+    def quantize(self):
+        node = self.node
+        assert node.op_type == "Where"
+        if not self.quantizer.force_quantize_no_input_check:
+            self.quantizer.new_nodes += [node]
+            return
+        (
+            data_found,
+            output_scale_name,
+            output_zp_name,
+            _,
+            _,
+        ) = self.quantizer._get_quantization_params(node.output[0])
+        (
+            q_input_names,
+            zero_point_names,
+            scale_names,
+            nodes,
+        ) = self.quantizer.quantize_activation(node, [1, 2])
+        if not data_found or q_input_names is None:
+            return super().quantize()
+        qlinear_output = node.output[0] + TENSOR_NAME_QUANT_SUFFIX
+        qlinear_output_name = node.name + "_quant" if node.name else ""
+
+        q_output = QuantizedValue(
+            node.output[0],
+            qlinear_output,
+            output_scale_name,
+            output_zp_name,
+            QuantizedValueType.Input,
+        )
+        self.quantizer.quantized_value_map[node.output[0]] = q_output
+
+        kwargs = {}
+        for attribute in node.attribute:
+            kwargs.update(attribute_to_kwarg(attribute))
+        kwargs["domain"] = ms_domain
+
+        qlwhere_inputs = [
+            node.input[0],
+            q_input_names[0],
+            scale_names[0],
+            zero_point_names[0],
+            q_input_names[1],
+            scale_names[1],
+            zero_point_names[1],
+            output_scale_name,
+            output_zp_name,
+        ]
+        qlwhere_node = onnx.helper.make_node(
+            "QLinearWhere", qlwhere_inputs, [qlinear_output], qlinear_output_name, **kwargs
+        )
+
+        self.quantizer.new_nodes += nodes
+        self.quantizer.new_nodes += [qlwhere_node]
+
+
+class QDQWhere(QDQOperatorBase):
+    def quantize(self):
+        node = self.node
+        assert node.op_type == "Where"
+        if self.quantizer.force_quantize_no_input_check:
+            if not self.quantizer.is_tensor_quantized(node.input[1]):
+                self.quantizer.quantize_activation_tensor(node.input[1])
+            if not self.quantizer.is_tensor_quantized(node.input[2]):
+                self.quantizer.quantize_activation_tensor(node.input[2])
+            if not self.disable_qdq_for_node_output:
+                for output in node.output:
+                    self.quantizer.quantize_activation_tensor(output)
+        elif (
+            self.quantizer.is_tensor_quantized(node.input[1])
+            and self.quantizer.is_tensor_quantized(node.input[2])
+            and not self.disable_qdq_for_node_output
+        ):
+            for output in node.output:
+                self.quantizer.quantize_activation_tensor(output)