fix: identify ConstantOfShape ONNX nodes as constant values

src/concrete/ml/pytest/torch_models.py

@@ -1564,3 +1564,35 @@ def forward(self, x):
         x = self.bn1(x)
         x = self.fc1(x)
         return x
+
+
+class ConstantOfShapeModel(nn.Module):
+    """Small model that uses a 'ConstantOfShape' ONNX node.
+
+    'ConstantOfShape' ONNX nodes appear when using torch operators like `zeros_like`. Most of the
+    time, ONNX seems to optimize the graph and remove them from it. This model is a specific case
+    where this optimization is not applied.
+    """
+
+    # pylint: disable-next=unused-argument
+    def __init__(self, input_output, activation_function, n_bits=2):
+        super().__init__()
+
+        self.quant = qnn.QuantIdentity(bit_width=n_bits)
+        self.conv = qnn.QuantConv2d(1, 1, 3, stride=1, bias=True, weight_bit_width=n_bits)
+
+    def forward(self, x):
+        """Forward pass.
+
+        Args:
+            x (torch.Tensor): The model's input.
+
+        Returns:
+            torch.Tensor: The model's output.
+
+        """
+        x = self.quant(x)
+        x = self.conv(x)
+        x = x + torch.zeros_like(x, dtype=torch.float32, requires_grad=True)
+
+        return x

src/concrete/ml/quantization/post_training.py

@@ -517,7 +517,11 @@ def _quantize_layers(self, *input_calibration_data: numpy.ndarray):
             assert_true(len(node.output) == 1)
 
             output_name = node.output[0]
-            if op_type == "Constant":
+
+            # 'ConstantOfShape' ONNX nodes appear when using torch operators like `zeros_like`.
+            # Most of the time, ONNX seems to optimize the graph and remove them from it. However,
+            # that is not always the case and we need to identify them as constant values as well
+            if op_type in ["Constant", "ConstantOfShape"]:
                 constant_values = ONNX_OPS_TO_NUMPY_IMPL["Constant"](**attributes)[0]
                 node_results[output_name] = constant_values
                 constants.add(output_name)

tests/torch/test_compile_torch.py

@@ -31,6 +31,7 @@
     CNNGrouped,
     CNNOther,
     ConcatFancyIndexing,
+    ConstantOfShapeModel,
     Conv1dModel,
     DoubleQuantQATMixNet,
     EncryptedMatrixMultiplicationModel,
@@ -625,7 +626,10 @@ def test_compile_torch_qat(
 
 @pytest.mark.parametrize(
     "model_class, input_output_feature, is_brevitas_qat",
-    [pytest.param(partial(MultiInputNNDifferentSize, is_brevitas_qat=True), [5, 10], True)],
+    [
+        pytest.param(partial(MultiInputNNDifferentSize, is_brevitas_qat=True), [5, 10], True),
+        pytest.param(ConstantOfShapeModel, (1, 8, 8), True),
+    ],
 )
 @pytest.mark.parametrize(
     "n_bits",