FIx: raw incremental linreg and revised test scope (#2816)

* FIx: raw incremental linreg and revised test scope

* oops

* oops 2

* lint

Author: ethanglaser

sklearnex/linear_model/incremental_linear.py

@@ -162,12 +162,12 @@ def _onedal_supported(self, method_name, *data):
     _onedal_gpu_supported = _onedal_supported
 
     def _onedal_predict(self, X, queue=None):
+        xp, _ = get_namespace(X)
+
         if not get_config()["use_raw_input"]:
             if sklearn_check_version("1.2"):
                 self._validate_params()
 
-            xp, _ = get_namespace(X)
-
             X = validate_data(
                 self,
                 X,

sklearnex/spmd/linear_model/tests/test_incremental_linear_spmd.py

@@ -118,10 +118,9 @@ def test_incremental_linear_regression_fit_spmd_gold(
 @pytest.mark.parametrize("num_blocks", [1, 2])
 @pytest.mark.parametrize("macro_block", [None, 1024])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
-@pytest.mark.parametrize("use_raw_input", [True, False])
 @pytest.mark.mpi
 def test_incremental_linear_regression_partial_fit_spmd_gold(
-    dataframe, queue, fit_intercept, num_blocks, macro_block, dtype, use_raw_input
+    dataframe, queue, fit_intercept, num_blocks, macro_block, dtype
 ):
     # Import spmd and non-SPMD algo
     from sklearnex.linear_model import IncrementalLinearRegression
@@ -179,9 +178,7 @@ def test_incremental_linear_regression_partial_fit_spmd_gold(
         local_dpt_y = _convert_to_dataframe(
             split_local_y[i], sycl_queue=queue, target_df=dataframe
         )
-        # Configure raw input status for spmd estimator
-        with config_context(use_raw_input=use_raw_input):
-            inclin_spmd.partial_fit(local_dpt_X, local_dpt_y)
+        inclin_spmd.partial_fit(local_dpt_X, local_dpt_y)
 
     inclin.fit(dpt_X, dpt_y)
 
@@ -272,6 +269,7 @@ def test_incremental_linear_regression_fit_spmd_random(
 @pytest.mark.parametrize("num_features", [5, 10])
 @pytest.mark.parametrize("macro_block", [None, 1024])
 @pytest.mark.parametrize("dtype", [np.float32, np.float64])
+@pytest.mark.parametrize("use_raw_input", [True, False])
 @pytest.mark.mpi
 def test_incremental_linear_regression_partial_fit_spmd_random(
     dataframe,
@@ -282,6 +280,7 @@ def test_incremental_linear_regression_partial_fit_spmd_random(
     num_features,
     macro_block,
     dtype,
+    use_raw_input,
 ):
     # Import spmd and non-SPMD algo
     from sklearnex.linear_model import IncrementalLinearRegression
@@ -328,7 +327,9 @@ def test_incremental_linear_regression_partial_fit_spmd_random(
         dpt_X = _convert_to_dataframe(X_split[i], sycl_queue=queue, target_df=dataframe)
         dpt_y = _convert_to_dataframe(y_split[i], sycl_queue=queue, target_df=dataframe)
 
-        inclin_spmd.partial_fit(local_dpt_X, local_dpt_y)
+        # Configure raw input status for spmd estimator
+        with config_context(use_raw_input=use_raw_input):
+            inclin_spmd.partial_fit(local_dpt_X, local_dpt_y)
         inclin.partial_fit(dpt_X, dpt_y)
 
     assert_allclose(_as_numpy(inclin.coef_), _as_numpy(inclin_spmd.coef_), atol=tol)
@@ -337,7 +338,8 @@ def test_incremental_linear_regression_partial_fit_spmd_random(
             _as_numpy(inclin.intercept_), _as_numpy(inclin_spmd.intercept_), atol=tol
         )
 
-    y_pred_spmd = inclin_spmd.predict(dpt_X_test)
+    with config_context(use_raw_input=use_raw_input):
+        y_pred_spmd = inclin_spmd.predict(dpt_X_test)
     y_pred = inclin.predict(dpt_X_test)
 
     assert_allclose(_as_numpy(y_pred_spmd), _as_numpy(y_pred), atol=tol)