Use scikit-learn's array-api to accelerate StandardScaler

docs/source/cuml-accel/limitations.rst

@@ -51,6 +51,11 @@ A few additional general notes:
   through 1.8. This ensures that cuML's implementation of scikit-learn
   compatible APIs works as expected.
 
+- Some estimators are accelerated using scikit-learn's experimental
+  `array-api`_ support. These estimators are only accelerated by ``cuml.accel``
+  when running with scikit-learn versions >= 1.8. Running with an older
+  version of scikit-learn will use an unaccelerated estimator.
+
 - Error and warning messages and formats may differ from scikit-learn. Some
   errors might present as C++ stacktraces instead of python errors.
 
@@ -424,10 +429,8 @@ StandardScaler
 
 ``StandardScaler`` will fall back to CPU in the following cases:
 
-- If ``partial_fit`` is called (incremental learning not supported on GPU).
-- If ``sample_weight`` is provided (weighted statistics not supported on GPU).
-- If ``X`` has object dtype, half precision (``float16``) dtype, or complex dtype (``complex64``, ``complex128``).
-- If ``X`` is a sparse matrix with integer dtype or in a format other than CSR or CSC.
+- If ``X`` is sparse
+- When run on scikit-learn < 1.8
 
 TargetEncoder
 ^^^^^^^^^^^^^
@@ -554,3 +557,4 @@ comparing the trustworthiness score (computed via
 
 
 .. _open an issue: https://github.com/rapidsai/cuml/issues
+.. _array-api: https://scikit-learn.org/stable/modules/array_api.html

python/cuml/cuml/accel/_overrides/sklearn/preprocessing.py

@@ -2,73 +2,17 @@
 # SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 # SPDX-License-Identifier: Apache-2.0
 #
-
-import cupyx.scipy.sparse as cupy_sparse
 import numpy as np
-from scipy import sparse as sp_sparse
 
 import cuml.preprocessing
-from cuml.accel.estimator_proxy import ProxyBase
+from cuml.accel.estimator_proxy import ArrayAPIProxyBase, ProxyBase
 from cuml.internals.interop import UnsupportedOnGPU
 
 __all__ = ("StandardScaler", "TargetEncoder")
 
 
-def _check_standardscaler_unsupported_inputs(X, **kwargs):
-    """Check if inputs are supported by cuML's StandardScaler on GPU.
-
-    Raises UnsupportedOnGPU for unsupported cases to trigger CPU fallback.
-    """
-    if kwargs.get("sample_weight") is not None:
-        raise UnsupportedOnGPU("sample_weight is not supported")
-
-    # Reject complex, object, and float16 dtypes
-    if hasattr(X, "dtype"):
-        if np.issubdtype(X.dtype, np.complexfloating):
-            raise UnsupportedOnGPU("complex dtype is not supported")
-        if X.dtype == np.object_:
-            raise UnsupportedOnGPU("object dtype is not supported")
-        if X.dtype == np.float16:
-            raise UnsupportedOnGPU("float16 dtype is not supported")
-
-    # Check for sparse matrices with unsupported properties
-    if sp_sparse.issparse(X):
-        if np.issubdtype(X.dtype, np.integer):
-            raise UnsupportedOnGPU(
-                "sparse matrix with integer dtype is not supported"
-            )
-        # cuML's StandardScaler algorithm only supports CSR/CSC formats.
-        if X.format not in ("csr", "csc"):
-            raise UnsupportedOnGPU(
-                f"sparse matrix format '{X.format}' is not supported"
-            )
-    elif cupy_sparse.issparse(X):
-        if np.issubdtype(X.dtype, np.integer):
-            raise UnsupportedOnGPU(
-                "sparse matrix with integer dtype is not supported"
-            )
-        # cuML's StandardScaler algorithm only supports CSR/CSC formats.
-        if X.format not in ("csr", "csc"):
-            raise UnsupportedOnGPU(
-                f"sparse matrix format '{X.format}' is not supported"
-            )
-
-
-class StandardScaler(ProxyBase):
-    _gpu_class = cuml.preprocessing.StandardScaler
-
-    def _gpu_fit(self, X, y=None, sample_weight=None):
-        kwargs = {"sample_weight": sample_weight}
-        _check_standardscaler_unsupported_inputs(X, **kwargs)
-        return self._gpu.fit(X, y)
-
-    def _gpu_fit_transform(self, X, y=None, **fit_params):
-        _check_standardscaler_unsupported_inputs(X, **fit_params)
-        return self._gpu.fit_transform(X, y, **fit_params)
-
-    def _gpu_partial_fit(self, X, y=None, sample_weight=None):
-        """partial_fit not supported on GPU - always fall back to CPU."""
-        raise UnsupportedOnGPU("partial_fit not supported on GPU")
+class StandardScaler(ArrayAPIProxyBase):
+    _cpu_class_path = "sklearn.preprocessing.StandardScaler"
 
 
 def _check_unsupported_inputs(X, y, cpu_model):

python/cuml/cuml/accel/_patches/sklearn/utils/_array_api.py

@@ -0,0 +1,22 @@
+# SPDX-FileCopyrightText: Copyright (c) 2026, NVIDIA CORPORATION.
+# SPDX-License-Identifier: Apache-2.0
+import functools
+
+from sklearn.utils._array_api import (
+    _check_array_api_dispatch as _orig_check_array_api_dispatch,
+)
+
+from cuml.internals.outputs import in_internal_context
+
+__all__ = ("_check_array_api_dispatch",)
+
+
+@functools.wraps(_orig_check_array_api_dispatch)
+def _check_array_api_dispatch(array_api_dispatch):
+    # sklearn's array-api support requires setting SCIPY_ARRAY_API=1, even
+    # though all uses we need it for don't rely on scipy. To work around this,
+    # we patch sklearn to disable the check when running within a cuml
+    # estimator. Usage outside of cuml estimators will still result in the
+    # proper error.
+    if not in_internal_context():
+        _orig_check_array_api_dispatch(array_api_dispatch)

python/cuml/cuml/accel/core.py

@@ -93,6 +93,7 @@ def debug(self, msg: str) -> None:
 _PATCHES = {
     "sklearn.pipeline",
     "sklearn.utils",
+    "sklearn.utils._array_api",
 }
 
 ACCELERATED_MODULES = sorted(_OVERRIDES.union(_PATCHES))