allow optional device selection

Author: matthiasdiener

arraycontext/impl/cupy/__init__.py

@@ -71,10 +71,14 @@ class CupyArrayContext(ArrayContext):
 
     _loopy_transform_cache: dict[lp.TranslationUnit, lp.ExecutorBase]
 
-    def __init__(self) -> None:
+    def __init__(self, device: int | None = None) -> None:
         super().__init__()
         self._loopy_transform_cache = {}
 
+        if device is not None:
+            import cupy as cp
+            cp.cuda.runtime.setDevice(device)
+
     array_types = (CupyNonObjectArray,)
 
     def _get_fake_numpy_namespace(self):