chore: also fix shape coherence with sklearn

src/concrete/ml/sklearn/base.py

@@ -694,6 +694,8 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
         Returns:
             numpy.ndarray: The post-processed predictions.
         """
+        assert isinstance(y_preds, numpy.ndarray), "Output predictions must be an array."
+
         return y_preds
 
 
@@ -805,8 +807,12 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
 
             # If the prediction array is 1D, transform the output into a 2D array [1-p, p],
             # with p the initial output probabilities
-            if y_preds.ndim == 1 or y_preds.shape[1] == 1:
-                y_preds = numpy.concatenate((1 - y_preds, y_preds), axis=1)
+            # This is similar to what is done in scikit-learn
+            if y_preds.ndim == 1:
+                return numpy.vstack([1 - y_preds, y_preds]).T
+
+            elif y_preds.shape[1] == 1:
+                return numpy.concatenate((1 - y_preds, y_preds), axis=1)
 
         # Else, apply the softmax operator
         else:
@@ -1387,8 +1393,13 @@ def quantize_input(self, X: numpy.ndarray) -> numpy.ndarray:
     def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray:
         self.check_model_is_fitted()
 
-        q_y_preds = self.output_quantizers[0].dequant(q_y_preds)
-        return q_y_preds
+        y_preds = self.output_quantizers[0].dequant(q_y_preds)
+
+        # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn
+        if y_preds.ndim == 2 and y_preds.shape[1] == 1:
+            return y_preds.ravel()
+
+        return y_preds
 
     def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]:
         assert self._tree_inference is not None, self._is_not_fitted_error_message()
@@ -1442,7 +1453,12 @@ def post_processing(self, y_preds: numpy.ndarray) -> numpy.ndarray:
         if not self._fhe_ensembling:
             y_preds = numpy.sum(y_preds, axis=-1)
 
-            assert_true(y_preds.ndim == 2, "y_preds should be a 2D array")
+            assert isinstance(y_preds, numpy.ndarray), "Output predictions must be an array."
+
+            # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn
+            if y_preds.ndim == 2 and y_preds.shape[1] == 1:
+                return y_preds.ravel()
+
             return y_preds
 
         return super().post_processing(y_preds)
@@ -1693,6 +1709,10 @@ def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray:
         # De-quantize the output values
         y_preds = self.output_quantizers[0].dequant(q_y_preds)
 
+        # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn
+        if y_preds.ndim == 2 and y_preds.shape[1] == 1:
+            return y_preds.ravel()
+
         return y_preds
 
     def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]:
@@ -1775,25 +1795,26 @@ def decision_function(
         """
         # Here, we want to use SklearnLinearModelMixin's `predict` method as confidence scores are
         # the dot product's output values, without any post-processing
-        y_preds = SklearnLinearModelMixin.predict(self, X, fhe=fhe)
-        return y_preds
+        y_scores = SklearnLinearModelMixin.predict(self, X, fhe=fhe)
+
+        return y_scores
 
     def predict_proba(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.ndarray:
-        y_logits = self.decision_function(X, fhe=fhe)
-        y_proba = self.post_processing(y_logits)
+        y_scores = self.decision_function(X, fhe=fhe)
+        y_proba = self.post_processing(y_scores)
         return y_proba
 
-    # In scikit-learn, the argmax is done on the logits directly, not the probabilities
+    # In scikit-learn, the argmax is done on the scores directly, not the probabilities
     def predict(self, X: Data, fhe: Union[FheMode, str] = FheMode.DISABLE) -> numpy.ndarray:
         # Compute the predicted scores
-        y_logits = self.decision_function(X, fhe=fhe)
+        y_scores = self.decision_function(X, fhe=fhe)
 
         # Retrieve the class with the highest score
         # If there is a single dimension, only compare the scores to 0
-        if y_logits.ndim == 1 or y_logits.shape[1] == 1:
-            y_preds = (y_logits > 0).astype(int)
+        if y_scores.ndim == 1:
+            y_preds = (y_scores > 0).astype(int)
         else:
-            y_preds = numpy.argmax(y_logits, axis=1)
+            y_preds = numpy.argmax(y_scores, axis=1)
 
         return self.classes_[y_preds]
 
@@ -2002,6 +2023,11 @@ def dequantize_output(self, q_y_preds: numpy.ndarray) -> numpy.ndarray:
         self.check_model_is_fitted()
         # We compute the sorted argmax in FHE, which are integers.
         # No need to de-quantize the output values
+
+        # If the preds have shape (n, 1), squeeze it to shape (n,) like in scikit-learn
+        if q_y_preds.ndim > 1 and q_y_preds.shape[1] == 1:
+            return q_y_preds.ravel()
+
         return q_y_preds
 
     def _get_module_to_compile(self) -> Union[Compiler, QuantizedModule]:

tests/sklearn/test_sklearn_models.py

@@ -204,13 +204,11 @@ def check_correctness_with_sklearn(
             y_scores_sklearn = sklearn_model.decision_function(x)
             y_scores_fhe = model.decision_function(x, fhe=fhe)
 
-            # Currently, for single target data sets, Concrete models' outputs have shape (n, 1)
-            # while scikit-learn models' outputs have shape (n, )
-            # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4029
-            # assert y_scores_sklearn.shape == y_scores_fhe.shape, (
-            #     "Method 'decision_function' outputs different shapes between scikit-learn and "
-            #     f"Concrete ML in FHE (fhe={fhe})"
-            # )
+            assert y_scores_sklearn.shape == y_scores_fhe.shape, (
+                "Method 'decision_function' outputs different shapes between scikit-learn and "
+                f"Concrete ML in FHE (fhe={fhe})"
+            )
+
             check_r2_score(y_scores_sklearn, y_scores_fhe, acceptance_score=acceptance_r2score)
 
         # LinearSVC models from scikit-learn do not provide a 'predict_proba' method
@@ -231,13 +229,10 @@ def check_correctness_with_sklearn(
     y_pred_sklearn = sklearn_model.predict(x)
     y_pred_fhe = model.predict(x, fhe=fhe)
 
-    # Currently, for single target data sets, Concrete models' outputs have shape (n, 1) while
-    # scikit-learn models' outputs have shape (n, )
-    # FIXME: https://github.com/zama-ai/concrete-ml-internal/issues/4029
-    # assert y_pred_sklearn.shape == y_pred_fhe.shape, (
-    #     "Method 'predict' outputs different shapes between scikit-learn and "
-    #     f"Concrete ML in FHE (fhe={fhe})"
-    # )
+    assert y_pred_sklearn.shape == y_pred_fhe.shape, (
+        "Method 'predict' outputs different shapes between scikit-learn and "
+        f"Concrete ML in FHE (fhe={fhe})"
+    )
 
     # If the model is a classifier, check that accuracies are similar
     if is_classifier_or_partial_classifier(model):
@@ -649,10 +644,14 @@ def check_separated_inference(model, fhe_circuit, x, check_float_array_equal):
         and get_model_name(model) != "KNeighborsClassifier"
     ):
         # For linear classifiers, the argmax is done on the scores directly, not the probabilities
+        # Also, it is handled differently if shape is (n,) instead of (n, 1)
         if is_model_class_in_a_list(model, _get_sklearn_linear_models()):
-            y_pred = numpy.argmax(y_scores, axis=-1)
+            if y_scores.ndim == 1:
+                y_pred = (y_scores > 0).astype(int)
+            else:
+                y_pred = numpy.argmax(y_scores, axis=1)
         else:
-            y_pred = numpy.argmax(y_pred, axis=-1)
+            y_pred = numpy.argmax(y_pred, axis=1)
 
         y_pred_class = model.predict(x, fhe="simulate")
 
@@ -1874,7 +1873,6 @@ def test_rounding_consistency_for_regular_models(
     n_bits,
     load_data,
     check_r2_score,
-    check_accuracy,
     is_weekly_option,
     verbose=True,
 ):
@@ -1894,7 +1892,7 @@ def test_rounding_consistency_for_regular_models(
     else:
         # Check `predict` for regressors
         predict_method = model.predict
-        metric = check_accuracy
+        metric = check_r2_score
 
     check_rounding_consistency(
         model,