Fix wrapping metaestimators in Pipeline in cuml.accel

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)

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)

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",

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)

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: