diff --git a/python/cuml/cuml/accel/_patches/sklearn/compose.py b/python/cuml/cuml/accel/_patches/sklearn/compose.py new file mode 100644 index 0000000000..8f8f9dec54 --- /dev/null +++ b/python/cuml/cuml/accel/_patches/sklearn/compose.py @@ -0,0 +1,25 @@ +# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION. +# SPDX-License-Identifier: Apache-2.0 +import functools + +from sklearn.compose import ColumnTransformer + +from cuml.internals.outputs import using_output_type + +__all__ = ("ColumnTransformer",) + + +def patch_method(name): + """Patch a ColumnTransformer method to ensure results returned as numpy.""" + orig_method = getattr(ColumnTransformer, name) + + @functools.wraps(orig_method) + def method(self, *args, **kwargs): + with using_output_type("numpy"): + return orig_method(self, *args, **kwargs) + + setattr(ColumnTransformer, name, method) + + +for method_name in ["fit", "fit_transform", "transform"]: + patch_method(method_name) diff --git a/python/cuml/cuml/accel/_patches/sklearn/pipeline.py b/python/cuml/cuml/accel/_patches/sklearn/pipeline.py index 629f0ea86a..a032f290a0 100644 --- a/python/cuml/cuml/accel/_patches/sklearn/pipeline.py +++ b/python/cuml/cuml/accel/_patches/sklearn/pipeline.py @@ -5,14 +5,14 @@ import cupy as cp from cupyx.scipy.sparse import issparse as is_cp_sparse -from sklearn.pipeline import Pipeline +from sklearn.pipeline import FeatureUnion, Pipeline from sklearn.utils.metaestimators import available_if from cuml.accel.estimator_proxy import is_proxy from cuml.internals.global_settings import GlobalSettings from cuml.internals.outputs import using_output_type -__all__ = ("Pipeline",) +__all__ = ("Pipeline", "FeatureUnion") def get_output_type(pipeline, reverse=False): @@ -54,7 +54,7 @@ def flat_steps(pipeline): return "cupy" -def patch_method(name): +def patch_pipeline_method(name): """Patch a sklearn Pipeline method to reduce device<->host transfers.""" orig_method = inspect.getattr_static(Pipeline, name) # Unwrap @available_if decorated methods @@ -105,4 +105,20 @@ def method(self, *args, **kwargs): "score_samples", "transform", ]: - patch_method(method_name) + patch_pipeline_method(method_name) + + +def patch_feature_union_method(name): + """Patch a FeatureUnion method to ensure results returned as numpy.""" + orig_method = getattr(FeatureUnion, name) + + @functools.wraps(orig_method) + def method(self, *args, **kwargs): + with using_output_type("numpy"): + return orig_method(self, *args, **kwargs) + + setattr(FeatureUnion, name, method) + + +for method_name in ["fit", "fit_transform", "transform"]: + patch_feature_union_method(method_name) diff --git a/python/cuml/cuml/accel/core.py b/python/cuml/cuml/accel/core.py index 22d15868e4..8ea280d054 100644 --- a/python/cuml/cuml/accel/core.py +++ b/python/cuml/cuml/accel/core.py @@ -92,6 +92,7 @@ def debug(self, msg: str) -> None: _PATCHES = { "sklearn.pipeline", + "sklearn.compose", "sklearn.utils", "sklearn.utils._array_api", "sklearn.utils.discovery", diff --git a/python/cuml/cuml_accel_tests/test_pipeline.py b/python/cuml/cuml_accel_tests/test_pipeline.py index cb0935d554..814b114ad2 100644 --- a/python/cuml/cuml_accel_tests/test_pipeline.py +++ b/python/cuml/cuml_accel_tests/test_pipeline.py @@ -13,6 +13,7 @@ from packaging.version import Version from sklearn.base import BaseEstimator from sklearn.cluster import DBSCAN, KMeans +from sklearn.compose import ColumnTransformer from sklearn.datasets import make_classification, make_regression from sklearn.decomposition import PCA, TruncatedSVD from sklearn.linear_model import ( @@ -28,8 +29,8 @@ KNeighborsRegressor, NearestNeighbors, ) -from sklearn.pipeline import Pipeline, make_pipeline -from sklearn.preprocessing import StandardScaler +from sklearn.pipeline import FeatureUnion, Pipeline, make_pipeline +from sklearn.preprocessing import RobustScaler, StandardScaler from umap import UMAP SKLEARN_18 = Version(sklearn.__version__) >= Version("1.8.0.dev0") @@ -356,3 +357,53 @@ def test_pipeline_classifier_predict_non_numeric_labels(patch_methods): assert isinstance(LogisticRegression.predict.args[0], cp.ndarray) # User-facing output is always numpy assert isinstance(out, np.ndarray) + + +@requires_sklearn_18 +def test_column_transfomer_in_pipeline_works(): + """Ensure outputs of steps in `ColumnTransformer` return as numpy""" + rng = np.random.default_rng(0) + X = rng.standard_normal((200, 20)).astype(np.float32) + y = rng.standard_normal(200).astype(np.float32) + + ct = ColumnTransformer( + [ + ("svd", TruncatedSVD(n_components=5), slice(0, 10)), + ("pass", "passthrough", slice(10, 20)), + ] + ) + + pipe = Pipeline( + [ + ("ct", ct), # Shouldn't be accelerated + ("scaler", RobustScaler()), # Not accelerated + ("ridge", Ridge()), # Accelerated + ] + ) + + pipe.fit(X, y) + + +@requires_sklearn_18 +def test_feature_union_in_pipeline_works(): + """Ensure outputs of steps in `FeatureUnion` return as numpy""" + rng = np.random.default_rng(0) + X = rng.standard_normal((200, 20)).astype(np.float32) + y = rng.standard_normal(200).astype(np.float32) + + union = FeatureUnion( + [ + ("svd", TruncatedSVD(n_components=2)), + ("pca", PCA(n_components=2)), + ] + ) + + pipe = Pipeline( + [ + ("features", union), # Shouldn't be accelerated + ("scaler", RobustScaler()), # Not accelerated + ("ridge", Ridge()), # Accelerated + ] + ) + + pipe.fit(X, y) diff --git a/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml b/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml index e19cee5cbf..0027602b97 100644 --- a/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml +++ b/python/cuml/cuml_accel_tests/upstream/scikit-learn/xfail-examples.yaml @@ -3,11 +3,6 @@ strict: false tests: - "decomposition::plot_faces_decomposition" -- reason: 'cuml.accel bug: implicit CuPy-to-NumPy conversion not handled in ColumnTransformer/Pipeline' - marker: cuml_accel_bugs - tests: - - "linear_model::plot_poisson_regression_non_normal_loss" - - "release_highlights::plot_release_highlights_1_1_0" - reason: 'cuml.accel bug: native crash in cuml PCA' marker: cuml_accel_bugs tests: