FIX: Return correct shape of probabilities from GPU logreg (#2645)

* fix shapes of prediction outputs

* more comprehensive testing for GPU cases

Author: david-cortes-intel

onedal/linear_model/logistic_regression.py

@@ -217,6 +217,7 @@ def _predict_proba(self, X):
         result = result = self._infer(X)
         _, xp, _ = _get_sycl_namespace(X)
         y = from_table(result.probabilities, like=X)
+        y = xp.reshape(y, -1)
         return xp.stack([1 - y, y], axis=1)
 
     def _predict_log_proba(self, X):

sklearnex/linear_model/tests/test_logreg.py

@@ -196,3 +196,25 @@ def logistic_model_function(predicted_probabilities, coefs):
             model_sklearnex.predict_proba(X), model_sklearnex.coef_
         )
         assert fn_sklearnex <= fn_sklearn
+
+
+@pytest.mark.parametrize(
+    "dataframe,queue", get_dataframes_and_queues(device_filter_="gpu")
+)
+def test_gpu_logreg_prediction_shapes(dataframe, queue):
+    if not queue or not queue.sycl_device.is_gpu:
+        pytest.skip("Test for GPU-only code branch")
+    from sklearnex.linear_model import LogisticRegression
+
+    X, y = make_classification(random_state=123)
+    X = _convert_to_dataframe(X, sycl_queue=queue, target_df=dataframe)
+    y = _convert_to_dataframe(y, sycl_queue=queue, target_df=dataframe)
+
+    model = LogisticRegression(solver="newton-cg").fit(X, y)
+    pred = model.predict(X)
+    pred_proba = model.predict_proba(X)
+    pred_log_proba = model.predict_log_proba(X)
+
+    np.testing.assert_array_equal(pred.shape, (X.shape[0],))
+    np.testing.assert_array_equal(pred_proba.shape, (X.shape[0], 2))
+    np.testing.assert_array_equal(pred_log_proba.shape, (X.shape[0], 2))