ci: validate PyTorch MPS support

.github/workflows/ci.yml

@@ -108,10 +108,57 @@ jobs:
           .smoke-venv/bin/python -c "import ml4t.models as m; print(m.__version__)"
           .smoke-venv/bin/python -c "from ml4t.models import PCAModel, SAEModel, StochasticDiscountFactorModel"
 
+  mps-smoke:
+    name: MPS Smoke
+    runs-on: macos-15
+    timeout-minutes: 10
+    steps:
+      - uses: actions/checkout@v6.0.3
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v8.2.0
+        with:
+          version: "latest"
+
+      - name: Set up Python
+        run: uv python install ${{ env.PYTHON_VERSION }}
+
+      - name: Install dependencies
+        run: uv sync --dev
+
+      - name: Verify PyTorch MPS
+        run: |
+          uv run python - <<'PY'
+          import platform
+
+          import torch
+
+          from ml4t.models._internal.torch_runtime import resolve_device, seed_torch
+
+          print(platform.platform())
+          print("machine:", platform.machine())
+          print("torch:", torch.__version__)
+          print("mps built:", torch.backends.mps.is_built())
+          print("mps available:", torch.backends.mps.is_available())
+
+          assert platform.machine() == "arm64"
+          assert torch.backends.mps.is_built()
+          assert torch.backends.mps.is_available()
+
+          device = resolve_device(torch, "mps")
+          assert device.type == "mps"
+
+          seed_torch(torch, 42, device)
+          x = torch.ones(4, device=device)
+          y = x * 2
+          assert y.device.type == "mps"
+          assert y.cpu().tolist() == [2.0, 2.0, 2.0, 2.0]
+          PY
+
   build:
     name: Build Package
     runs-on: ubuntu-latest
-    needs: [lint, typecheck, test, test-package-smoke]
+    needs: [lint, typecheck, test, test-package-smoke, mps-smoke]
     steps:
       - uses: actions/checkout@v6.0.3
         with:

src/ml4t/models/_internal/torch_runtime.py

@@ -16,12 +16,21 @@ def import_torch() -> Any:
 
 
 def resolve_device(torch: Any, requested: str) -> Any:
-    if requested.startswith("cuda") and torch.cuda.is_available():
+    requested_device = requested.strip().lower()
+    if requested_device.startswith("cuda") and torch.cuda.is_available():
         return torch.device(requested)
+    mps_backend = getattr(torch.backends, "mps", None)
+    if requested_device == "mps" and mps_backend is not None and mps_backend.is_available():
+        return torch.device("mps")
     return torch.device("cpu")
 
 
 def seed_torch(torch: Any, seed: int, device: Any) -> None:
     torch.manual_seed(seed)
-    if getattr(device, "type", "cpu") == "cuda":
+    device_type = getattr(device, "type", "cpu")
+    if device_type == "cuda":
         torch.cuda.manual_seed_all(seed)
+    elif device_type == "mps":
+        mps_manual_seed = getattr(getattr(torch, "mps", None), "manual_seed", None)
+        if callable(mps_manual_seed):
+            mps_manual_seed(seed)

tests/test_torch_runtime.py

@@ -0,0 +1,94 @@
+from __future__ import annotations
+
+from dataclasses import dataclass
+
+from ml4t.models._internal.torch_runtime import resolve_device, seed_torch
+
+
+@dataclass(frozen=True)
+class _FakeDevice:
+    type: str
+
+
+class _FakeCuda:
+    def __init__(self, available: bool) -> None:
+        self.available = available
+        self.seeds: list[int] = []
+
+    def is_available(self) -> bool:
+        return self.available
+
+    def manual_seed_all(self, seed: int) -> None:
+        self.seeds.append(seed)
+
+
+class _FakeMPSBackend:
+    def __init__(self, available: bool) -> None:
+        self.available = available
+
+    def is_available(self) -> bool:
+        return self.available
+
+
+class _FakeMPS:
+    def __init__(self) -> None:
+        self.seeds: list[int] = []
+
+    def manual_seed(self, seed: int) -> None:
+        self.seeds.append(seed)
+
+
+class _FakeTorch:
+    def __init__(self, *, cuda_available: bool = False, mps_available: bool = False) -> None:
+        self.cuda = _FakeCuda(cuda_available)
+        self.backends = type("Backends", (), {"mps": _FakeMPSBackend(mps_available)})()
+        self.mps = _FakeMPS()
+        self.seeds: list[int] = []
+
+    def device(self, requested: str) -> _FakeDevice:
+        return _FakeDevice(requested.split(":", maxsplit=1)[0])
+
+    def manual_seed(self, seed: int) -> None:
+        self.seeds.append(seed)
+
+
+def test_resolve_device_cpu() -> None:
+    torch = _FakeTorch()
+
+    assert resolve_device(torch, "cpu").type == "cpu"
+
+
+def test_resolve_device_cuda_when_available() -> None:
+    torch = _FakeTorch(cuda_available=True)
+
+    assert resolve_device(torch, "cuda:0").type == "cuda"
+
+
+def test_resolve_device_cuda_falls_back_to_cpu() -> None:
+    torch = _FakeTorch(cuda_available=False)
+
+    assert resolve_device(torch, "cuda:0").type == "cpu"
+
+
+def test_resolve_device_mps_when_available() -> None:
+    torch = _FakeTorch(mps_available=True)
+
+    assert resolve_device(torch, "mps").type == "mps"
+
+
+def test_resolve_device_mps_falls_back_to_cpu() -> None:
+    torch = _FakeTorch(mps_available=False)
+
+    assert resolve_device(torch, "mps").type == "cpu"
+
+
+def test_seed_torch_dispatches_by_device() -> None:
+    torch = _FakeTorch(cuda_available=True, mps_available=True)
+
+    seed_torch(torch, 11, _FakeDevice("cpu"))
+    seed_torch(torch, 13, _FakeDevice("cuda"))
+    seed_torch(torch, 17, _FakeDevice("mps"))
+
+    assert torch.seeds == [11, 13, 17]
+    assert torch.cuda.seeds == [13]
+    assert torch.mps.seeds == [17]