scikit-learn · ogrisel · Aug 18, 2023 · Aug 18, 2023 · Aug 18, 2023 · Aug 18, 2023
diff --git a/doc/modules/array_api.rst b/doc/modules/array_api.rst
@@ -118,4 +118,30 @@ To run these checks you need to install
 test environment. To run the full set of checks you need to install both
 `PyTorch <https://pytorch.org/>`_ and `CuPy <https://cupy.dev/>`_ and have
 a GPU. Checks that can not be executed or have missing dependencies will be
-automatically skipped.
+automatically skipped. Therefore it's important to run the tests with the
+`-v` flag to see which checks are skipped:
+
+.. prompt:: bash $
+
+    $ pip install array-api-compat  # and other libraries as needed
+    $ pytest -k "array_api" -v
+
+Note on MPS device support
+--------------------------
+
+On macOS, PyTorch can use the Metal Performance Shaders (MPS) to access
+hardware accelerators (e.g. the internal GPU component of the M1 or M2 chips).
+However, the MPS device support for PyTorch is incomplete at the time of
+writing. See the following github issue for more details:
+
+- https://github.com/pytorch/pytorch/issues/77764
+
+To enable the MPS support in PyTorch, set the environment variable
+`PYTORCH_ENABLE_MPS_FALLBACK=1` before running the tests:
+
+.. prompt:: bash $
+
+    $ PYTORCH_ENABLE_MPS_FALLBACK=1 pytest -k "array_api" -v
+
+At the time of writing all scikit-learn tests should pass, however, the
+computational speed is not necessarily better than with the CPU device.
diff --git a/doc/whats_new/v1.4.rst b/doc/whats_new/v1.4.rst
@@ -107,8 +107,8 @@ Changelog
 - |Enhancement| :class:`decomposition.PCA` now supports the Array API for the
   `full` and `randomized` solvers (with QR power iterations). See
   :ref:`array_api` for more details.
-  :pr:`26315` by :user:`Mateusz Sokół <mtsokol>` and
-  :user:`Olivier Grisel <ogrisel>`.
+  :pr:`26315` and :pr:`27098` by :user:`Mateusz Sokół <mtsokol>`,
+  :user:`Olivier Grisel <ogrisel>` and :user:` Edoardo Abati <EdAbati>`.
 
 :mod:`sklearn.ensemble`
 .......................

diff --git a/sklearn/decomposition/_base.py b/sklearn/decomposition/_base.py
@@ -14,7 +14,7 @@
 from scipy import linalg
 
 from ..base import BaseEstimator, ClassNamePrefixFeaturesOutMixin, TransformerMixin
-from ..utils._array_api import _add_to_diagonal, get_namespace
+from ..utils._array_api import _add_to_diagonal, device, get_namespace
 from ..utils.validation import check_is_fitted
 
 
@@ -47,7 +47,9 @@ def get_covariance(self):
             components_ = components_ * xp.sqrt(exp_var[:, np.newaxis])
         exp_var_diff = exp_var - self.noise_variance_
         exp_var_diff = xp.where(
-            exp_var > self.noise_variance_, exp_var_diff, xp.asarray(0.0)
+            exp_var > self.noise_variance_,
+            exp_var_diff,
+            xp.asarray(0.0, device=device(exp_var)),
         )
         cov = (components_.T * exp_var_diff) @ components_
         _add_to_diagonal(cov, self.noise_variance_, xp)
@@ -87,7 +89,9 @@ def get_precision(self):
             components_ = components_ * xp.sqrt(exp_var[:, np.newaxis])
         exp_var_diff = exp_var - self.noise_variance_
         exp_var_diff = xp.where(
-            exp_var > self.noise_variance_, exp_var_diff, xp.asarray(0.0)
+            exp_var > self.noise_variance_,
+            exp_var_diff,
+            xp.asarray(0.0, device=device(exp_var)),
         )
         precision = components_ @ components_.T / self.noise_variance_
         _add_to_diagonal(precision, 1.0 / exp_var_diff, xp)

diff --git a/sklearn/utils/_array_api.py b/sklearn/utils/_array_api.py
@@ -45,6 +45,7 @@ def yield_namespace_device_dtype_combinations():
                 ("cpu", "cuda"), ("float64", "float32")
             ):
                 yield array_namespace, device, dtype
+            yield array_namespace, "mps", "float32"
         else:
             yield array_namespace, None, None
 
@@ -442,7 +443,9 @@ def _weighted_sum(sample_score, sample_weight, normalize=False, xp=None):
         return float(numpy.average(sample_score_np, weights=sample_weight_np))
 
     if not xp.isdtype(sample_score.dtype, "real floating"):
-        sample_score = xp.astype(sample_score, xp.float64)
+        # We move to cpu device ahead of time since certain devices may not support
+        # float64, but we want the same precision for all devices and namespaces.
+        sample_score = xp.astype(xp.asarray(sample_score, device="cpu"), xp.float64)
 
     if sample_weight is not None:
         sample_weight = xp.asarray(sample_weight)

diff --git a/sklearn/utils/estimator_checks.py b/sklearn/utils/estimator_checks.py
@@ -868,9 +868,19 @@ def _array_api_for_tests(array_namespace, device, dtype):
     # This is because `cupy` is not the same as the compatibility wrapped
     # namespace of a CuPy array.
     xp = array_api_compat.get_namespace(array_mod.asarray(1))
-
     if array_namespace == "torch" and device == "cuda" and not xp.has_cuda:
         raise SkipTest("PyTorch test requires cuda, which is not available")
+    elif array_namespace == "torch" and device == "mps" and not xp.has_mps:
+        if not xp.backends.mps.is_built():
+            raise SkipTest(
+                "MPS is not available because the current PyTorch install was not "
+                "built with MPS enabled."
+            )
+        else:
+            raise SkipTest(
+                "MPS is not available because the current MacOS version is not 12.3+ "
+                "and/or you do not have an MPS-enabled device on this machine."
+            )
     elif array_namespace in {"cupy", "cupy.array_api"}:  # pragma: nocover
         import cupy