Move MPS device check to cheetah.utils

cheetah/utils/__init__.py

@@ -1,2 +1,3 @@
+from .device import is_mps_available_and_functional  # noqa: F401
 from .kde import kde_histogram_1d, kde_histogram_2d  # noqa: F401
 from .unique_name_generator import UniqueNameGenerator  # noqa: F401

cheetah/utils/device.py

@@ -0,0 +1,13 @@
+import torch
+
+
+def is_mps_available_and_functional():
+    """Check if MPS is available and functional (for GitHub Actions)."""
+    if not torch.backends.mps.is_available():
+        return False
+    try:
+        # Try to allocate a small tensor on the MPS device
+        torch.tensor([1.0], device="mps")
+        return True
+    except RuntimeError:
+        return False

tests/__init__.py

@@ -1,13 +0,0 @@
-import torch
-
-
-def is_mps_available_and_functional():
-    """Check if MPS is available and functional (for GitHub Actions)."""
-    if not torch.backends.mps.is_available():
-        return False
-    try:
-        # Try to allocate a small tensor on the MPS device
-        torch.tensor([1.0], device="mps")
-        return True
-    except RuntimeError:
-        return False

tests/test_bmad_conversion.py

@@ -2,8 +2,7 @@
 import torch
 
 import cheetah
-
-from . import is_mps_available_and_functional
+from cheetah.utils import is_mps_available_and_functional
 
 
 def test_bmad_tutorial():

tests/test_device_dtype.py

@@ -2,8 +2,7 @@
 import torch
 
 import cheetah
-
-from . import is_mps_available_and_functional
+from cheetah.utils import is_mps_available_and_functional
 
 
 @pytest.mark.parametrize(