MNT refactor name of device function to array_device (scikit-learn#34543)

Author: OmarManzoor

sklearn/_loss/tests/test_loss.py

@@ -37,12 +37,10 @@
 from sklearn.utils import assert_all_finite
 from sklearn.utils._array_api import (
     _atol_for_type,
+    array_device,
     move_to,
     yield_namespace_device_dtype_combinations,
 )
-from sklearn.utils._array_api import (
-    device as array_api_device,
-)
 from sklearn.utils._testing import (
     _array_api_for_tests,
     create_memmap_backed_data,
@@ -1421,7 +1419,7 @@ def _assert_array_api_result(
             move_to(result_xp, xp=np, device="cpu"), result_np, rtol=rtol, atol=atol
         )
         assert result_xp.dtype == raw_prediction_xp.dtype
-        assert array_api_device(result_xp) == array_api_device(raw_prediction_xp)
+        assert array_device(result_xp) == array_device(raw_prediction_xp)
 
     if method_name == "gradient_proba" and loss_class != HalfMultinomialLoss:
         # `gradient_proba` is only valid for HalfMultinomialLoss

sklearn/calibration.py

@@ -410,9 +410,9 @@ def fit(self, X, y, sample_weight=None, **fit_params):
                 {"splitter": {}, "estimator": {"fit": fit_kwargs}}
             )
 
-        xp, is_array_api, device_ = get_namespace_and_device(X)
+        xp, is_array_api, device = get_namespace_and_device(X)
         if is_array_api:
-            y, sample_weight = move_to(y, sample_weight, xp=xp, device=device_)
+            y, sample_weight = move_to(y, sample_weight, xp=xp, device=device)
         # Check that each cross-validation fold can have at least one
         # example per class
         if isinstance(self.cv, int):
@@ -527,8 +527,8 @@ def predict_proba(self, X):
         check_is_fitted(self)
         # Compute the arithmetic mean of the predictions of the calibrated
         # classifiers
-        xp, _, device_ = get_namespace_and_device(X)
-        mean_proba = xp.zeros((_num_samples(X), self.classes_.shape[0]), device=device_)
+        xp, _, device = get_namespace_and_device(X)
+        mean_proba = xp.zeros((_num_samples(X), self.classes_.shape[0]), device=device)
         for calibrated_classifier in self.calibrated_classifiers_:
             proba = calibrated_classifier.predict_proba(X)
             mean_proba += proba
@@ -981,7 +981,7 @@ def _convert_to_logits(decision_values, eps=1e-12, xp=None):
     -------
     logits : ndarray of shape (n_samples, n_classes)
     """
-    xp, _, device_ = get_namespace_and_device(decision_values, xp=xp)
+    xp, _, device = get_namespace_and_device(decision_values, xp=xp)
     decision_values = check_array(
         decision_values, dtype=[xp.float64, xp.float32], ensure_2d=False
     )
@@ -993,7 +993,7 @@ def _convert_to_logits(decision_values, eps=1e-12, xp=None):
         row_sums_to_one = xp.all(
             xpx.isclose(
                 xp.sum(decision_values, axis=1),
-                xp.asarray(1.0, device=device_, dtype=decision_values.dtype),
+                xp.asarray(1.0, device=device, dtype=decision_values.dtype),
             )
         )
 

sklearn/covariance/_empirical_covariance.py

@@ -272,10 +272,10 @@ def fit(self, X, y=None):
         self : object
             Returns the instance itself.
         """
-        xp, _, device_ = get_namespace_and_device(X)
-        X = validate_data(self, X, dtype=supported_float_dtypes(xp, device_))
+        xp, _, device = get_namespace_and_device(X)
+        X = validate_data(self, X, dtype=supported_float_dtypes(xp, device))
         if self.assume_centered:
-            self.location_ = xp.zeros(X.shape[1], dtype=X.dtype, device=device_)
+            self.location_ = xp.zeros(X.shape[1], dtype=X.dtype, device=device)
         else:
             self.location_ = xp.mean(X, axis=0)
         covariance = empirical_covariance(X, assume_centered=self.assume_centered)
@@ -388,9 +388,9 @@ def mahalanobis(self, X):
         dist : ndarray of shape (n_samples,)
             Squared Mahalanobis distances of the observations.
         """
-        xp, _, device_ = get_namespace_and_device(X)
+        xp, _, device = get_namespace_and_device(X)
         X = validate_data(
-            self, X, reset=False, dtype=supported_float_dtypes(xp, device_)
+            self, X, reset=False, dtype=supported_float_dtypes(xp, device)
         )
 
         precision = self.get_precision()

sklearn/covariance/_shrunk_covariance.py

@@ -623,10 +623,10 @@ def fit(self, X, y=None):
         """
         # Not calling the parent object to fit, to avoid computing the
         # covariance matrix (and potentially the precision)
-        xp, _, device_ = get_namespace_and_device(X)
-        X = validate_data(self, X, dtype=supported_float_dtypes(xp, device_))
+        xp, _, device = get_namespace_and_device(X)
+        X = validate_data(self, X, dtype=supported_float_dtypes(xp, device))
         if self.assume_centered:
-            self.location_ = xp.zeros(X.shape[1], dtype=X.dtype, device=device_)
+            self.location_ = xp.zeros(X.shape[1], dtype=X.dtype, device=device)
         else:
             self.location_ = xp.mean(X, axis=0)
         covariance, shrinkage = _ledoit_wolf(

sklearn/decomposition/_base.py

@@ -13,7 +13,7 @@
     ClassNamePrefixFeaturesOutMixin,
     TransformerMixin,
 )
-from sklearn.utils._array_api import _add_to_diagonal, device, get_namespace
+from sklearn.utils._array_api import _add_to_diagonal, array_device, get_namespace
 from sklearn.utils.validation import check_array, check_is_fitted, validate_data
 
 
@@ -48,7 +48,7 @@ def get_covariance(self):
         exp_var_diff = xp.where(
             exp_var > self.noise_variance_,
             exp_var_diff,
-            xp.asarray(0.0, device=device(exp_var), dtype=exp_var.dtype),
+            xp.asarray(0.0, device=array_device(exp_var), dtype=exp_var.dtype),
         )
         cov = (components_.T * exp_var_diff) @ components_
         _add_to_diagonal(cov, self.noise_variance_, xp)
@@ -90,7 +90,7 @@ def get_precision(self):
         exp_var_diff = xp.where(
             exp_var > self.noise_variance_,
             exp_var_diff,
-            xp.asarray(0.0, device=device(exp_var)),
+            xp.asarray(0.0, device=array_device(exp_var)),
         )
         precision = components_ @ components_.T / self.noise_variance_
         _add_to_diagonal(precision, 1.0 / exp_var_diff, xp)

sklearn/decomposition/_pca.py

@@ -15,7 +15,7 @@
 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 _cov, device, get_namespace
+from sklearn.utils._array_api import _cov, array_device, get_namespace
 from sklearn.utils._param_validation import Interval, RealNotInt, StrOptions
 from sklearn.utils.extmath import _randomized_svd, fast_logdet, svd_flip
 from sklearn.utils.sparsefuncs import _implicit_column_offset, mean_variance_axis
@@ -641,7 +641,7 @@ def _fit_full(self, X, n_components, xp, is_array_api_compliant):
             n_components = (
                 xp.searchsorted(
                     ratio_cumsum,
-                    xp.asarray(n_components, device=device(ratio_cumsum)),
+                    xp.asarray(n_components, device=array_device(ratio_cumsum)),
                     side="right",
                 )
                 + 1

sklearn/decomposition/tests/test_pca.py

@@ -13,10 +13,10 @@
 from sklearn.decomposition._pca import _assess_dimension, _infer_dimension
 from sklearn.utils._array_api import (
     _atol_for_type,
+    array_device,
     move_to,
     yield_namespace_device_dtype_combinations,
 )
-from sklearn.utils._array_api import device as array_device
 from sklearn.utils._test_common.instance_generator import _get_check_estimator_ids
 from sklearn.utils._testing import _array_api_for_tests, assert_allclose
 from sklearn.utils.estimator_checks import (

sklearn/discriminant_analysis.py

@@ -22,8 +22,8 @@
 from sklearn.preprocessing import StandardScaler
 from sklearn.utils._array_api import (
     _expit,
+    array_device,
     check_same_namespace,
-    device,
     get_namespace,
     size,
 )
@@ -113,7 +113,9 @@ def _class_means(X, y):
     """
     xp, is_array_api_compliant = get_namespace(X)
     classes, y = xp.unique_inverse(y)
-    means = xp.zeros((classes.shape[0], X.shape[1]), device=device(X), dtype=X.dtype)
+    means = xp.zeros(
+        (classes.shape[0], X.shape[1]), device=array_device(X), dtype=X.dtype
+    )
 
     if is_array_api_compliant:
         for i in range(classes.shape[0]):
@@ -602,7 +604,9 @@ def _solve_svd(self, X, y):
         std = xp.std(Xc, axis=0)
         # avoid division by zero in normalization
         std[std == 0] = 1.0
-        fac = xp.asarray(1.0 / (n_samples - n_classes), dtype=X.dtype, device=device(X))
+        fac = xp.asarray(
+            1.0 / (n_samples - n_classes), dtype=X.dtype, device=array_device(X)
+        )
 
         # 2) Within variance scaling
         X = xp.sqrt(fac) * (Xc / std)

sklearn/linear_model/_base.py

@@ -164,7 +164,7 @@ def _preprocess_data(
     sample_weight_sqrt : ndarray of shape (n_samples, ) or None
         `np.sqrt(sample_weight)`
     """
-    xp, _, device_ = get_namespace_and_device(X, y, sample_weight)
+    xp, _, device = get_namespace_and_device(X, y, sample_weight)
     n_samples, n_features = X.shape
     X_is_sparse = sp.issparse(X)
 
@@ -173,7 +173,7 @@ def _preprocess_data(
             X,
             copy=copy,
             accept_sparse=["csr", "csc"],
-            dtype=supported_float_dtypes(xp, device=device_),
+            dtype=supported_float_dtypes(xp, device=device),
         )
         y = check_array(y, dtype=X.dtype, copy=True, ensure_2d=False)
     else:
@@ -198,15 +198,15 @@ def _preprocess_data(
         y_offset = _average(y, axis=0, weights=sample_weight, xp=xp)
         y -= y_offset
     else:
-        X_offset = xp.zeros(n_features, dtype=X.dtype, device=device_)
+        X_offset = xp.zeros(n_features, dtype=X.dtype, device=device)
         if y.ndim == 1:
-            y_offset = xp.asarray(0.0, dtype=dtype_, device=device_)
+            y_offset = xp.asarray(0.0, dtype=dtype_, device=device)
         else:
-            y_offset = xp.zeros(y.shape[1], dtype=dtype_, device=device_)
+            y_offset = xp.zeros(y.shape[1], dtype=dtype_, device=device)
 
     # X_scale is no longer needed. It is a historic artifact from the
     # time where linear model exposed the normalize parameter.
-    X_scale = xp.ones(n_features, dtype=X.dtype, device=device_)
+    X_scale = xp.ones(n_features, dtype=X.dtype, device=device)
 
     if sample_weight is not None and rescale_with_sw:
         # Sample weight can be implemented via a simple rescaling.
@@ -410,7 +410,7 @@ def predict(self, X):
             Vector containing the class labels for each sample.
         """
         check_same_namespace(X, self, attribute="coef_", method="predict")
-        xp, _, device_ = get_namespace_and_device(X)
+        xp, _, device = get_namespace_and_device(X)
         scores = self.decision_function(X)
         if len(scores.shape) == 1:
             indices = xp.astype(scores > 0, indexing_dtype(xp))
@@ -424,7 +424,7 @@ def predict(self, X):
         if isinstance(y_pred[0], str):
             return y_pred
         else:
-            return move_to(y_pred, xp=xp, device=device_)
+            return move_to(y_pred, xp=xp, device=device)
 
     def _predict_proba_lr(self, X):
         """Probability estimation for OvR logistic regression.

sklearn/linear_model/_glm/glm.py

@@ -262,7 +262,7 @@ def fit(self, X, y, sample_weight=None):
                 f"l1_ratio={self.l1_ratio}."
             )
             raise ValueError(msg)
-        xp, _, device_ = get_namespace_and_device(X)
+        xp, _, device = get_namespace_and_device(X)
         X, y = validate_data(
             self,
             X,
@@ -280,10 +280,10 @@ def fit(self, X, y, sample_weight=None):
             # losses.
             sample_weight = _check_sample_weight(sample_weight, X, dtype=loss_dtype)
 
-        y, sample_weight = move_to(y, sample_weight, xp=xp, device=device_)
+        y, sample_weight = move_to(y, sample_weight, xp=xp, device=device)
 
         n_samples, n_features = X.shape
-        self._base_loss = self._get_loss(xp=xp, device=device_)
+        self._base_loss = self._get_loss(xp=xp, device=device)
 
         linear_loss = LinearModelLoss(
             base_loss=self._base_loss,
@@ -362,7 +362,7 @@ def fit(self, X, y, sample_weight=None):
             coef = xp.asarray(
                 coef.copy(order="C" if not _is_numpy_namespace(xp) else "K"),
                 dtype=X.dtype,
-                device=device_,
+                device=device,
             )
         elif self.solver == "newton-cg":
             func = linear_loss.loss
@@ -514,8 +514,8 @@ def score(self, X, y, sample_weight=None):
             # losses.
             sample_weight = _check_sample_weight(sample_weight, X, dtype=y.dtype)
 
-        xp, _, device_ = get_namespace_and_device(X)
-        y, sample_weight = move_to(y, sample_weight, xp=xp, device=device_)
+        xp, _, device = get_namespace_and_device(X)
+        y, sample_weight = move_to(y, sample_weight, xp=xp, device=device)
 
         base_loss = self._base_loss