ENH use xp.cumulative_sum and xp.searchsorted directly instead of stable_cumsum (#31994)

otizonaizit · web-flow · commit 450cb20733c4 · 2025-08-22T22:06:19.000+05:00
diff --git a/sklearn/decomposition/_pca.py b/sklearn/decomposition/_pca.py
@@ -15,9 +15,9 @@
 from sklearn.decomposition._base import _BasePCA
 from sklearn.utils import check_random_state
 from sklearn.utils._arpack import _init_arpack_v0
-from sklearn.utils._array_api import _convert_to_numpy, get_namespace
+from sklearn.utils._array_api import device, get_namespace
 from sklearn.utils._param_validation import Interval, RealNotInt, StrOptions
-from sklearn.utils.extmath import _randomized_svd, fast_logdet, stable_cumsum, svd_flip
+from sklearn.utils.extmath import _randomized_svd, fast_logdet, svd_flip
 from sklearn.utils.sparsefuncs import _implicit_column_offset, mean_variance_axis
 from sklearn.utils.validation import check_is_fitted, validate_data
 
@@ -655,23 +655,15 @@ def _fit_full(self, X, n_components, xp, is_array_api_compliant):
             # side='right' ensures that number of features selected
             # their variance is always greater than n_components float
             # passed. More discussion in issue: #15669
-            if is_array_api_compliant:
-                # Convert to numpy as xp.cumsum and xp.searchsorted are not
-                # part of the Array API standard yet:
-                #
-                # https://github.com/data-apis/array-api/issues/597
-                # https://github.com/data-apis/array-api/issues/688
-                #
-                # Furthermore, it's not always safe to call them for namespaces
-                # that already implement them: for instance as
-                # cupy.searchsorted does not accept a float as second argument.
-                explained_variance_ratio_np = _convert_to_numpy(
-                    explained_variance_ratio_, xp=xp
+            ratio_cumsum = xp.cumulative_sum(explained_variance_ratio_)
+            n_components = (
+                xp.searchsorted(
+                    ratio_cumsum,
+                    xp.asarray(n_components, device=device(ratio_cumsum)),
+                    side="right",
                 )
-            else:
-                explained_variance_ratio_np = explained_variance_ratio_
-            ratio_cumsum = stable_cumsum(explained_variance_ratio_np)
-            n_components = np.searchsorted(ratio_cumsum, n_components, side="right") + 1
+                + 1
+            )
 
         # Compute noise covariance using Probabilistic PCA model
         # The sigma2 maximum likelihood (cf. eq. 12.46)
