FIX Explicitly disable array API dispatching for reference computation (scikit-learn#34378)

Author: betatim

sklearn/_loss/tests/test_loss.py

@@ -1455,9 +1455,12 @@ def _assert_array_api_result(
 
     method = getattr(loss_instance, method_name)
     array_api_method = getattr(array_api_loss_instance, method_name)
-    result_np = method(
-        y_true=y_true, raw_prediction=raw_prediction, sample_weight=sample_weight_np
-    )
+    with config_context(array_api_dispatch=False):
+        result_np = method(
+            y_true=y_true,
+            raw_prediction=raw_prediction,
+            sample_weight=sample_weight_np,
+        )
     with config_context(array_api_dispatch=True):
         result_xp = array_api_method(
             y_true=y_true_xp,

sklearn/covariance/tests/test_covariance.py

@@ -395,7 +395,8 @@ def test_empirical_covariance_array_api(
     X_np = X.astype(dtype_name, copy=False)
     X_xp = xp.asarray(X_np, device=device)
 
-    result_np = empirical_covariance(X_np, assume_centered=assume_centered)
+    with config_context(array_api_dispatch=False):
+        result_np = empirical_covariance(X_np, assume_centered=assume_centered)
 
     with config_context(array_api_dispatch=True):
         result_xp = empirical_covariance(X_xp, assume_centered=assume_centered)
@@ -418,7 +419,8 @@ def test_ledoit_wolf_shrinkage_array_api(
     X_np = X[:, :n_features].astype(dtype_name, copy=False)
     X_xp = xp.asarray(X_np, device=device)
 
-    shrinkage_np = ledoit_wolf_shrinkage(X_np)
+    with config_context(array_api_dispatch=False):
+        shrinkage_np = ledoit_wolf_shrinkage(X_np)
 
     with config_context(array_api_dispatch=True):
         shrinkage_xp = ledoit_wolf_shrinkage(X_xp)
@@ -436,10 +438,11 @@ def test_log_likelihood_array_api(array_namespace, device_name, dtype_name):
     xp, device = _array_api_for_tests(array_namespace, device_name, dtype_name)
 
     X_np = X.astype(dtype_name, copy=False)
-    emp_cov_np = empirical_covariance(X_np)
-    est = EmpiricalCovariance().fit(X_np)
-    precision_np = est.get_precision()
-    result_np = log_likelihood(emp_cov_np, precision_np)
+    with config_context(array_api_dispatch=False):
+        emp_cov_np = empirical_covariance(X_np)
+        est = EmpiricalCovariance().fit(X_np)
+        precision_np = est.get_precision()
+        result_np = log_likelihood(emp_cov_np, precision_np)
 
     # Downcast to target dtype before converting because `empirical_covariance` can
     # return float64 for numpy and MPS only supports float32
@@ -469,8 +472,9 @@ def test_score_array_api(array_namespace, device_name, dtype_name, store_precisi
         lw_xp = LedoitWolf(store_precision=store_precision).fit(X_xp)
         score_xp = lw_xp.score(X_xp)
 
-    lw_np = LedoitWolf(store_precision=store_precision).fit(X_np)
-    score_np = lw_np.score(X_np)
+    with config_context(array_api_dispatch=False):
+        lw_np = LedoitWolf(store_precision=store_precision).fit(X_np)
+        score_np = lw_np.score(X_np)
 
     assert isinstance(score_xp, float)
     assert_allclose(score_np, score_xp, atol=_atol_for_type(dtype_name))
@@ -488,8 +492,9 @@ def test_error_norm_array_api(array_namespace, device_name, dtype_name, norm):
     X_np = X.astype(dtype_name, copy=False)
     X_xp = xp.asarray(X_np, device=device)
 
-    lw_np = LedoitWolf().fit(X_np)
-    comp_cov_np = empirical_covariance(X_np)
+    with config_context(array_api_dispatch=False):
+        lw_np = LedoitWolf().fit(X_np)
+        comp_cov_np = empirical_covariance(X_np)
 
     with config_context(array_api_dispatch=True):
         lw_xp = LedoitWolf().fit(X_xp)
@@ -498,7 +503,8 @@ def test_error_norm_array_api(array_namespace, device_name, dtype_name, norm):
         comp_cov_xp = xp.asarray(comp_cov_np.astype(dtype_name), device=device)
         result_xp = lw_xp.error_norm(comp_cov_xp, norm=norm)
 
-    result_np = lw_np.error_norm(comp_cov_np, norm=norm)
+    with config_context(array_api_dispatch=False):
+        result_np = lw_np.error_norm(comp_cov_np, norm=norm)
     assert_allclose(result_np, float(result_xp), atol=_atol_for_type(dtype_name))
 
 
@@ -513,8 +519,9 @@ def test_mahalanobis_array_api(array_namespace, device_name, dtype_name):
     X_np = X.astype(dtype_name, copy=False)
     X_xp = xp.asarray(X_np, device=device)
 
-    lw_np = LedoitWolf().fit(X_np)
-    dist_np = lw_np.mahalanobis(X_np)
+    with config_context(array_api_dispatch=False):
+        lw_np = LedoitWolf().fit(X_np)
+        dist_np = lw_np.mahalanobis(X_np)
 
     with config_context(array_api_dispatch=True):
         lw_xp = LedoitWolf().fit(X_xp)

sklearn/decomposition/tests/test_pca.py

@@ -942,9 +942,10 @@ def check_array_api_get_precision(
     iris_np = iris.data.astype(dtype_name)
     iris_xp = xp.asarray(iris_np, device=device)
 
-    estimator.fit(iris_np)
-    precision_np = estimator.get_precision()
-    covariance_np = estimator.get_covariance()
+    with config_context(array_api_dispatch=False):
+        estimator.fit(iris_np)
+        precision_np = estimator.get_precision()
+        covariance_np = estimator.get_covariance()
 
     rtol = 2e-4 if iris_np.dtype == "float32" else 2e-7
     with config_context(array_api_dispatch=True):
@@ -1057,10 +1058,11 @@ def test_pca_mle_array_api_compliance(
     X_xp = xp.asarray(X, device=device)
     y_xp = xp.asarray(y, device=device)
 
-    est.fit(X, y)
+    with config_context(array_api_dispatch=False):
+        est.fit(X, y)
 
-    components_np = est.components_
-    explained_variance_np = est.explained_variance_
+        components_np = est.components_
+        explained_variance_np = est.explained_variance_
 
     est_xp = clone(est)
     with config_context(array_api_dispatch=True):

sklearn/linear_model/_glm/tests/test_glm.py

@@ -1282,13 +1282,15 @@ def test_poisson_regressor_array_api_compliance(
 
     params = dict(alpha=1, solver="lbfgs", max_iter=500)
     params["tol"] = 3e-6 if dtype_name == "float32" else 1e-13
-    glm_np = PoissonRegressor(**params).fit(X_np, y_np, sample_weight=sample_weight)
-    assert glm_np.n_iter_ < glm_np.max_iter
+    with config_context(array_api_dispatch=False):
+        glm_np = PoissonRegressor(**params).fit(X_np, y_np, sample_weight=sample_weight)
+        assert glm_np.n_iter_ < glm_np.max_iter
 
-    # Test that alpha was not too large for meaningful testing.
-    assert np.abs(glm_np.coef_).max() > 0.1
+        # Test that alpha was not too large for meaningful testing.
+        assert np.abs(glm_np.coef_).max() > 0.1
+
+        predict_np = glm_np.predict(X_np)
 
-    predict_np = glm_np.predict(X_np)
     atol = _atol_for_type(dtype_name)
     rtol = 2e-3 if dtype_name == "float32" else 3e-7
 

sklearn/linear_model/tests/test_logistic.py

@@ -2878,16 +2878,23 @@ def test_logistic_regression_array_api_compliance(
     lr_params = dict(C=1e-2, solver=solver, max_iter=500, class_weight=class_weight)
     lr_params["tol"] = 1e-6 if dtype_name == "float32" else 1e-12
 
-    lr_np = LogisticRegression(**lr_params).fit(X_np, y_np, sample_weight=sample_weight)
-    # Make sure that the reference fit converged.
-    assert lr_np.n_iter_ < lr_np.max_iter
+    # Compute the reference fit and predictions with array API dispatch explicitly
+    # disabled so that they always follow the NumPy code path, independently of
+    # the default value of the `array_api_dispatch` config flag.
+    with config_context(array_api_dispatch=False):
+        lr_np = LogisticRegression(**lr_params).fit(
+            X_np, y_np, sample_weight=sample_weight
+        )
+        # Make sure that the reference fit converged.
+        assert lr_np.n_iter_ < lr_np.max_iter
 
-    # Test that C was not too large for meaningful testing.
-    assert np.abs(lr_np.coef_).max() > 0.1
+        # Test that C was not too large for meaningful testing.
+        assert np.abs(lr_np.coef_).max() > 0.1
+
+        predict_proba_np = lr_np.predict_proba(X_np)
+        preditct_log_proba_np = lr_np.predict_log_proba(X_np)
+        prediction_np = lr_np.predict(X_np)
 
-    predict_proba_np = lr_np.predict_proba(X_np)
-    preditct_log_proba_np = lr_np.predict_log_proba(X_np)
-    prediction_np = lr_np.predict(X_np)
     if solver == "lbfgs":
         atol = _atol_for_type(dtype_name) * 10
         rtol = 1e-3 if dtype_name == "float32" else 1e-7

sklearn/linear_model/tests/test_ridge.py

@@ -1460,9 +1460,10 @@ def check_array_api_attributes(
     X_iris_xp = xp.asarray(X_iris_np, device=device)
     y_iris_xp = xp.asarray(y_iris_np, device=device)
 
-    estimator.fit(X_iris_np, y_iris_np)
-    coef_np = estimator.coef_
-    intercept_np = estimator.intercept_
+    with config_context(array_api_dispatch=False):
+        estimator.fit(X_iris_np, y_iris_np)
+        coef_np = estimator.coef_
+        intercept_np = estimator.intercept_
 
     with config_context(array_api_dispatch=True):
         estimator_xp = clone(estimator).fit(X_iris_xp, y_iris_xp)
@@ -1535,9 +1536,11 @@ def test_ridge_classifier_multilabel_array_api(
     X, y = make_multilabel_classification(random_state=0)
     X_np = X.astype(dtype_name)
     y_np = y.astype(dtype_name)
-    ridge_np = estimator.fit(X_np, y_np)
-    pred_np = ridge_np.predict(X_np)
-    classes_np = ridge_np.classes_.copy()
+    with config_context(array_api_dispatch=False):
+        ridge_np = estimator.fit(X_np, y_np)
+        pred_np = ridge_np.predict(X_np)
+        classes_np = ridge_np.classes_.copy()
+
     with config_context(array_api_dispatch=True):
         X_xp, y_xp = xp.asarray(X_np, device=device), xp.asarray(y_np, device=device)
         ridge_xp = estimator.fit(X_xp, y_xp)
@@ -1562,8 +1565,9 @@ def test_ridge_per_target_alpha_array_api(array_namespace, device_name, dtype_na
     y_np = y.astype(dtype_name)
     alphas = np.asarray([1e-2, 0.1, 1.0], dtype=dtype_name)
 
-    ridge_np = Ridge(alpha=alphas, solver="svd").fit(X_np, y_np)
-    pred_np = ridge_np.predict(X_np)
+    with config_context(array_api_dispatch=False):
+        ridge_np = Ridge(alpha=alphas, solver="svd").fit(X_np, y_np)
+        pred_np = ridge_np.predict(X_np)
 
     with config_context(array_api_dispatch=True):
         X_xp, y_xp = xp.asarray(X_np, device=device), xp.asarray(y_np, device=device)

sklearn/mixture/tests/test_gaussian_mixture.py

@@ -1521,7 +1521,8 @@ def test_gaussian_mixture_array_api_compliance(
         init_params=init_params,
         **additional_kwargs,
     )
-    gmm.fit(X)
+    with sklearn.config_context(array_api_dispatch=False):
+        gmm.fit(X)
 
     X_xp = xp.asarray(X, device=device)
 
@@ -1588,26 +1589,33 @@ def test_gaussian_mixture_array_api_compliance(
     )
 
     # Check methods
+    with sklearn.config_context(array_api_dispatch=False):
+        predict_np = gmm.predict(X)
+        predict_proba_np = gmm.predict_proba(X)
+        score_samples_np = gmm.score_samples(X)
+        score_np = gmm.score(X)
+        aic_np = gmm.aic(X)
+        bic_np = gmm.bic(X)
+        sample_X, sample_y = gmm.sample(10)
+
     assert (
-        adjusted_rand_score(gmm.predict(X), move_to(predict_xp, xp=np, device="cpu"))
-        > 0.95
+        adjusted_rand_score(predict_np, move_to(predict_xp, xp=np, device="cpu")) > 0.95
     )
     assert_allclose(
-        gmm.predict_proba(X),
+        predict_proba_np,
         move_to(predict_proba_xp, xp=np, device="cpu"),
         rtol=increased_rtol,
         atol=increased_atol,
     )
     assert_allclose(
-        gmm.score_samples(X),
+        score_samples_np,
         move_to(score_samples_xp, xp=np, device="cpu"),
         rtol=increased_rtol,
     )
     # comparing Python float so need explicit rtol when X has dtype float32
-    assert_allclose(gmm.score(X), score_xp, rtol=default_rtol)
-    assert_allclose(gmm.aic(X), aic_xp, rtol=default_rtol)
-    assert_allclose(gmm.bic(X), bic_xp, rtol=default_rtol)
-    sample_X, sample_y = gmm.sample(10)
+    assert_allclose(score_np, score_xp, rtol=default_rtol)
+    assert_allclose(aic_np, aic_xp, rtol=default_rtol)
+    assert_allclose(bic_np, bic_xp, rtol=default_rtol)
     # generated samples are float64 so need explicit rtol when X has dtype float32
     assert_allclose(
         sample_X, move_to(sample_X_xp, xp=np, device="cpu"), rtol=default_rtol

sklearn/model_selection/tests/test_split.py

@@ -1365,9 +1365,11 @@ def test_array_api_train_test_split(
     y_np = y.astype(dtype_name)
     y_xp = xp.asarray(y_np, device=device)
 
-    X_train_np, X_test_np, y_train_np, y_test_np = train_test_split(
-        X_np, y, random_state=0, shuffle=shuffle, stratify=stratify
-    )
+    with config_context(array_api_dispatch=False):
+        X_train_np, X_test_np, y_train_np, y_test_np = train_test_split(
+            X_np, y, random_state=0, shuffle=shuffle, stratify=stratify
+        )
+
     with config_context(array_api_dispatch=True):
         if stratify is not None:
             stratify_xp = xp.asarray(stratify)

sklearn/model_selection/tests/test_validation.py

@@ -2729,7 +2729,8 @@ def test_cross_val_predict_array_api_compliance(
     with config_context(array_api_dispatch=True):
         pred_xp = cross_val_predict(estimator, X_xp, y_xp, cv=cv)
 
-    pred_np = cross_val_predict(estimator, X_np, y_np, cv=cv)
+    with config_context(array_api_dispatch=False):
+        pred_np = cross_val_predict(estimator, X_np, y_np, cv=cv)
     assert_allclose(
         move_to(pred_xp, xp=np, device="cpu"), pred_np, atol=_atol_for_type(dtype_name)
     )

sklearn/preprocessing/tests/test_data.py

@@ -189,11 +189,12 @@ def test_standard_scaler_sample_weight_array_api(
     yw = np.ones(Xw.shape[0]).astype(dtype_name, copy=False)
     X_test = np.array([[1.5, 2.5, 3.5], [3.5, 4.5, 5.5]]).astype(dtype_name, copy=False)
 
-    scaler = StandardScaler()
-    scaler.fit(X, y)
+    with config_context(array_api_dispatch=False):
+        scaler = StandardScaler()
+        scaler.fit(X, y)
 
-    scaler_w = StandardScaler()
-    scaler_w.fit(Xw, yw, sample_weight=sample_weight)
+        scaler_w = StandardScaler()
+        scaler_w.fit(Xw, yw, sample_weight=sample_weight)
 
     # Test array-api support and correctness.
     X_xp = xp.asarray(X, device=device)
@@ -2159,7 +2160,8 @@ def test_binarizer_array_api_int(array_namespace, device_name, dtype_name):
     for dtype_name_ in [dtype_name, "int32", "int64"]:
         X_np = np.reshape(np.asarray([0, 1, 2, 3, 4], dtype=dtype_name_), (-1, 1))
         X_xp = xp.asarray(X_np, device=device)
-        binarized_np = Binarizer(threshold=2.5).fit_transform(X_np)
+        with config_context(array_api_dispatch=False):
+            binarized_np = Binarizer(threshold=2.5).fit_transform(X_np)
         with config_context(array_api_dispatch=True):
             binarized_xp = Binarizer(threshold=2.5).fit_transform(X_xp)
         assert_array_equal(move_to(binarized_xp, xp=np, device="cpu"), binarized_np)