Add CuPy support to ArrayStore via Array API

.github/workflows/testing.yml

@@ -65,7 +65,7 @@ jobs:
     strategy:
       max-parallel: 12 # All in parallel.
       matrix:
-        os: [ubuntu-latest, macos-latest, windows-latest]
+        os: [ubuntu-latest]
         python-version: ["3.12"]
     runs-on: ${{ matrix.os }}
     steps:
@@ -77,6 +77,8 @@ jobs:
         run: pip install .[dev]
       - name: Install torch
         run: pip install torch
+      - name: Install cupy
+        run: pip install cupy-cuda12x
       - name: Test core
         run: >
           pytest tests/archives tests/emitters tests/schedulers

HISTORY.md

@@ -6,7 +6,7 @@
 
 #### API
 
-- Support array backends via Python array API Standard ({pr}`573`, {pr}`571`)
+- Support array backends via Python array API Standard ({issue}`570`)
 - **Backwards-incompatible:** Remove raw_dict methods from ArrayStore
   ({pr}`575`)
 

README.md

@@ -124,10 +124,11 @@ python -c "import ribs; print(ribs.__version__)"
 You should see a version number in the output.
 
 **Experimental:** Pyribs is experimenting with adding support for running QD
-algorithms in PyTorch via the
+algorithms in PyTorch and CuPy via the
 [Python array API standard](https://data-apis.org/array-api/latest/). To enable
-this functionality, [install PyTorch](https://pytorch.org), such as with
-`pip install torch`.
+this functionality, install [PyTorch](https://pytorch.org), such as with
+`pip install torch`, and [CuPy](https://cupy.dev), such as with
+`pip install cupy-cuda12x`.
 
 ## Usage
 

ribs/archives/_array_store.py

@@ -4,7 +4,13 @@
 from enum import IntEnum
 from functools import cached_property
 
-from array_api_compat import is_numpy_array, is_numpy_namespace, is_torch_array
+from array_api_compat import (is_cupy_array, is_numpy_array, is_numpy_namespace,
+                              is_torch_array)
+
+try:
+    from array_api_compat import cupy as cp
+except ImportError:
+    pass
 
 from ribs._utils import arr_readonly, xp_namespace
 from ribs.archives._archive_data_frame import ArchiveDataFrame
@@ -292,10 +298,12 @@ def _convert_to_numpy(arr):
             return arr
         elif is_torch_array(arr):
             return arr.cpu().detach().numpy()
+        elif is_cupy_array(arr):
+            return cp.asnumpy(arr)
         else:
             raise NotImplementedError(
                 "The pandas return type is currently only supported "
-                "with numpy and torch arrays.")
+                "with NumPy, PyTorch, and CuPy arrays.")
 
     def retrieve(self, indices, fields=None, return_type="dict"):
         """Collects data at the given indices.

tests/conftest.py

@@ -24,6 +24,21 @@
 except ImportError:
     pass
 
+try:
+    import cupy as cp
+    from cupy_backends.cuda.api.runtime import CUDARuntimeError
+
+    cp.empty(0)  # Triggers CUDARuntimeError if there is no GPU available.
+
+    xp_available_backends.append(
+        pytest.param((cp, cp.cuda.Device(0)), id="cupy-gpu"))
+except ImportError:
+    # CuPy not installed.
+    pass
+except CUDARuntimeError:
+    # GPU not available.
+    pass
+
 
 @pytest.fixture(params=xp_available_backends)
 def xp_and_device(request):