From 0bd1c4d44820a829fdba6f168c91095e1a44de34 Mon Sep 17 00:00:00 2001 From: MalekAl-Abed Date: Mon, 13 Jul 2026 18:03:26 +0200 Subject: [PATCH 1/2] Add PyTorch MPS device support --- pyRadPlan/core/xp_utils/__init__.py | 17 +++++++++++++---- test/core/xp_utils/test_initialization.py | 13 +++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/pyRadPlan/core/xp_utils/__init__.py b/pyRadPlan/core/xp_utils/__init__.py index 4e23cf0..5941c06 100644 --- a/pyRadPlan/core/xp_utils/__init__.py +++ b/pyRadPlan/core/xp_utils/__init__.py @@ -134,15 +134,24 @@ def choose_device(xp: Optional[ArrayNamespace] = None, gpu_index: int = 0) -> Un xp : Optional[ArrayNamespace], optional The array namespace. If None, the preferred backend is used. gpu_index : int, optional - GPU device index to use when a GPU backend is selected. Default is 0. - Pass ``1``, ``2``, etc. for multi-GPU setups. + GPU device index used for CUDA backends. Ignored when the selected backend is MPS. + Default is 0. Pass ``1``, ``2``, etc. for multi-GPU setups. """ if xp is None: xp = choose_array_api_namespace() if PREFER_GPU: - if array_api_compat.is_torch_namespace(xp) and pytorch_gpu_available(): - return f"cuda:{gpu_index}" + if array_api_compat.is_torch_namespace(xp): + if ( + torch is not None + and hasattr(torch.backends, "mps") + and torch.backends.mps.is_available() + ): + return "mps" + + elif pytorch_gpu_available(): + return f"cuda:{gpu_index}" + if array_api_compat.is_cupy_namespace(xp) and cupy_available(): return str(gpu_index) diff --git a/test/core/xp_utils/test_initialization.py b/test/core/xp_utils/test_initialization.py index b5bf09d..79eb72e 100644 --- a/test/core/xp_utils/test_initialization.py +++ b/test/core/xp_utils/test_initialization.py @@ -137,6 +137,19 @@ def test_choose_device_torch_multi_gpu(): assert choose_device(xp, gpu_index=1) == "cuda:1" +@pytest.mark.skipif( + not (HAS_TORCH and hasattr(torch.backends, "mps")), + reason="PyTorch MPS backend unavailable", +) +def test_choose_device_torch_mps(monkeypatch): + import array_api_compat.torch as xp + + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(torch.cuda, "is_available", lambda: False) + + assert choose_device(xp) == "mps" + + @pytest.mark.skipif(not (HAS_CUPY and CUPY_CUDA_AVAILABLE), reason="CuPy GPU not available") def test_choose_device_cupy(): """Test choose_device with cupy namespace.""" From e13077f2e86b4899f98b206e55fad1c7f8b6f387 Mon Sep 17 00:00:00 2001 From: MalekAl-Abed Date: Mon, 13 Jul 2026 18:12:56 +0200 Subject: [PATCH 2/2] Add PyTorch MPS device support -- Changelog --- CHANGELOG.md | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/CHANGELOG.md b/CHANGELOG.md index ead6d94..eb2a6d1 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -9,9 +9,13 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ## [0.4.1] - 2026-06-16 + ### Added - multiple corresponding units for given quantities in the GUI +- support for Apple Metal Performance Shaders (MPS) devices when using the PyTorch backend + + ### Fixed