TST fix estimator checks when set_output is called on the instance (scikit-learn#29869)

adrinjalali · web-flow · commit 13de540b3aa7 · 2024-09-23T12:03:49.000-04:00
diff --git a/doc/whats_new/v1.6.rst b/doc/whats_new/v1.6.rst
@@ -367,6 +367,11 @@ Changelog
   calling :func:`utils.validation.check_non_negative`.
   :pr:`29540` by :user:`Tamara Atanasoska <tamaraatanasoska>`.
 
+- |FIX| :func:`utils.estimator_checks.parametrize_with_checks` and
+  :func:`utils.estimator_checks.check_estimator` now support estimators that
+  have `set_output` called on them.
+  :pr:`29869` by `Adrin Jalali`_.
+
 - |Enhancement| :func:`utils.validation.check_is_fitted` now passes on stateless
   estimators. An estimator can indicate it's stateless by setting the `requires_fit`
   tag. See :ref:`estimator_tags` for more information.
diff --git a/sklearn/utils/estimator_checks.py b/sklearn/utils/estimator_checks.py
@@ -40,6 +40,7 @@
 from ..model_selection._validation import _safe_split
 from ..pipeline import make_pipeline
 from ..preprocessing import StandardScaler, scale
+from ..utils import _safe_indexing
 from ..utils._array_api import (
     _atol_for_type,
     _convert_to_numpy,
@@ -1458,8 +1459,8 @@ def check_methods_sample_order_invariance(name, estimator_orig):
 
         if hasattr(estimator, method):
             assert_allclose_dense_sparse(
-                getattr(estimator, method)(X)[idx],
-                getattr(estimator, method)(X[idx]),
+                _safe_indexing(getattr(estimator, method)(X), idx),
+                getattr(estimator, method)(_safe_indexing(X, idx)),
                 atol=1e-9,
                 err_msg=msg,
             )
@@ -1804,6 +1805,9 @@ def check_estimators_dtypes(name, estimator_orig):
 def check_transformer_preserve_dtypes(name, transformer_orig):
     # check that dtype are preserved meaning if input X is of some dtype
     # X_transformed should be from the same dtype.
+    transformer = clone(transformer_orig)
+    if hasattr(transformer, "set_output"):
+        transformer.set_output(transform="default")
     X, y = make_blobs(
         n_samples=30,
         centers=[[0, 0, 0], [1, 1, 1]],
@@ -1815,7 +1819,6 @@ def check_transformer_preserve_dtypes(name, transformer_orig):
 
     for dtype in get_tags(transformer_orig).transformer_tags.preserves_dtype:
         X_cast = X.astype(dtype)
-        transformer = clone(transformer_orig)
         set_random_state(transformer)
         X_trans1 = transformer.fit_transform(X_cast, y)
         X_trans2 = transformer.fit(X_cast, y).transform(X_cast)
@@ -3777,7 +3780,9 @@ def check_fit_idempotent(name, estimator_orig):
     for method in check_methods:
         if hasattr(estimator, method):
             new_result = getattr(estimator, method)(X_test)
-            if np.issubdtype(new_result.dtype, np.floating):
+            if hasattr(new_result, "dtype") and np.issubdtype(
+                new_result.dtype, np.floating
+            ):
                 tol = 2 * np.finfo(new_result.dtype).eps
             else:
                 tol = 2 * np.finfo(np.float64).eps
diff --git a/sklearn/utils/tests/test_estimator_checks.py b/sklearn/utils/tests/test_estimator_checks.py
@@ -27,6 +27,7 @@
 )
 from sklearn.mixture import GaussianMixture
 from sklearn.neighbors import KNeighborsRegressor
+from sklearn.preprocessing import StandardScaler
 from sklearn.svm import SVC, NuSVC
 from sklearn.utils import _array_api, all_estimators, deprecated
 from sklearn.utils._param_validation import Interval, StrOptions
@@ -1417,3 +1418,16 @@ def _more_tags(self):
     # to exist so that third party estimators can easily support multiple sklearn
     # versions.
     check_estimator_tags_renamed("OkayEstimator", OkayEstimator())
+
+
+# Test that set_output doesn't make the tests to fail.
+def test_estimator_with_set_output():
+    # Doing this since pytest is not available for this file.
+    for lib in ["pandas", "polars"]:
+        try:
+            importlib.__import__(lib)
+        except ImportError:
+            raise SkipTest(f"Library {lib} is not installed")
+
+        estimator = StandardScaler().set_output(transform=lib)
+        check_estimator(estimator)
