FIX skip array API tests when running with device="mps" without the PYTORCH_ENABLE_MPS_FALLBACK env var (#27199)

betatim · ogrisel · web-flow · commit b7d80a0317d0 · 2023-09-07T12:23:35.000-04:00
Co-authored-by: Olivier Grisel &lt;olivier.grisel@ensta.org&gt;
diff --git a/sklearn/utils/_testing.py b/sklearn/utils/_testing.py
@@ -1071,17 +1071,26 @@ def _array_api_for_tests(array_namespace, device, dtype):
     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():
+    elif array_namespace == "torch" and device == "mps":
+        if os.getenv("PYTORCH_ENABLE_MPS_FALLBACK") != "1":
+            # For now we need PYTORCH_ENABLE_MPS_FALLBACK=1 for all estimators to work
+            # when using the MPS device.
             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."
+                "Skipping MPS device test because PYTORCH_ENABLE_MPS_FALLBACK is not "
+                "set."
             )
+        if 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
 
