Simplify conftest

Author: crusaderky

tests/conftest.py

@@ -2,7 +2,7 @@
 
 from collections.abc import Callable, Generator
 from contextlib import suppress
-from functools import cache, partial, wraps
+from functools import partial, wraps
 from types import ModuleType
 from typing import ParamSpec, TypeVar, cast
 
@@ -102,28 +102,6 @@ def wrapper(*args: P.args, **kwargs: P.kwargs) -> T:  # numpydoc ignore=GL08
         return wrapper
 
 
-@cache
-def _jax_cuda_device() -> Device | None:
-    """Return a CUDA device for JAX, if available."""
-    import jax
-
-    try:
-        return jax.devices("cuda")[0]
-    except Exception:
-        return None
-
-
-@cache
-def _torch_cuda_device() -> Device | None:
-    """Return a CUDA device for PyTorch, if available."""
-    import torch
-
-    try:
-        return torch.empty((0,), device=torch.device("cuda")).device
-    except Exception:
-        return None
-
-
 @pytest.fixture
 def xp(
     library: Backend, request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch
@@ -168,15 +146,18 @@ def xp(
         jax.config.update("jax_enable_x64", True)  # type: ignore[no-untyped-call,unused-ignore]
 
         if library == Backend.JAX_GPU:
-            device = _jax_cuda_device()
-            if device is None:
+            try:
+                device = jax.devices("cuda")[0]
+            except RuntimeError:
                 pytest.skip("no cuda device available")
         else:
             device = jax.devices("cpu")[0]
         jax.config.update("jax_default_device", device)
 
     elif library == Backend.TORCH_GPU:
-        if _torch_cuda_device() is None:
+        import torch.cuda
+
+        if not torch.cuda.is_available():
             pytest.skip("no cuda device available")
         xp.set_default_device("cuda")