Start Array API infrastructure (#573)

## Description

<!-- Provide a brief description of the PR's purpose here. -->

Towards #570

## TODO

<!-- Notable points that this PR has either accomplished or will
accomplish. -->

- [x] Add `array-api-compat` to deps -- no version specified because it
seems like a fairly flexible library
- [x] Add instructions for installing PyTorch -- it is now marked as
experimental in the README
- [x] Add CI for testing with PyTorch -- it is added as a separate test
in `ci.yaml` that copies the existing tests -- eventually this should be
merged into the main tests (#574)
- [x] Set up testing for Array API libraries with `xp_and_device`
fixture in `tests/conftest.py`
- [x] Update testing instructions in `tests/README.md`

## Status

- [x] I have read the guidelines in

[CONTRIBUTING.md](https://github.com/icaros-usc/pyribs/blob/master/CONTRIBUTING.md)
- [x] I have formatted my code using `yapf`
- [x] I have tested my code by running `pytest`
- [x] I have linted my code with `pylint`
- [x] I have added a one-line description of my change to the changelog
in
      `HISTORY.md`
- [x] This PR is ready to go

Author: btjanaka

.github/workflows/testing.yml

@@ -61,6 +61,37 @@ jobs:
         run: pip install pymoo
       - name: Test pymoo extra
         run: pytest tests/emitters_pymoo tests/schedulers_pymoo
+  test_array_api:
+    strategy:
+      max-parallel: 12 # All in parallel.
+      matrix:
+        os: [ubuntu-latest, macos-latest, windows-latest]
+        python-version: ["3.12"]
+    runs-on: ${{ matrix.os }}
+    steps:
+      - uses: actions/checkout@v4
+      - uses: conda-incubator/setup-miniconda@v3
+        with:
+          python-version: ${{ matrix.python-version }}
+      - name: Install core deps
+        run: pip install .[dev]
+      - name: Install torch
+        run: pip install torch
+      - name: Test core
+        run: >
+          pytest tests/archives tests/emitters tests/schedulers
+      - name: Install visualize dep
+        run: pip install .[visualize]
+      - name: Test visualize extra
+        run: pytest tests/visualize
+      - name: Install cma
+        run: pip install cma
+      - name: Test cma extra
+        run: pytest tests/emitters_pycma
+      - name: Install pymoo
+        run: pip install pymoo
+      - name: Test pymoo extra
+        run: pytest tests/emitters_pymoo tests/schedulers_pymoo
   visualize_qdax:
     runs-on: ubuntu-latest
     steps:

HISTORY.md

@@ -1,5 +1,13 @@
 # History
 
+## (Forthcoming)
+
+### Changelog
+
+#### API
+
+- Support array backends via Python array API Standard ({pr}`573`)
+
 ## 0.8.1
 
 This release makes some minor updates to 0.8.0.

README.md

@@ -95,7 +95,7 @@ in pyribs operates as follows:
 
 ## Installation
 
-pyribs supports Python 3.9 and above. The vast majority of users can install
+Pyribs supports Python 3.9 and above. The vast majority of users can install
 pyribs by running:
 
 ```bash
@@ -123,6 +123,12 @@ 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
+[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`.
+
 ## Usage
 
 Here we show an example application of CMA-ME in pyribs. To initialize the

pyproject.toml

@@ -29,15 +29,14 @@ classifiers = [
 ]
 requires-python = ">=3.9.0"
 dependencies = [
-  # numpy>=1.17.0 is when default_rng becomes available;
-  # scikit-learn 1.1.0 requires numpy 1.17.3+
-  "numpy>=1.17.3",
-  "numpy_groupies>=0.9.16",  # Supports Python 3.7 and up.
+  "array_api_compat",
   "numba>=0.51.0",
+  "numpy>=1.22.0", # numpy>=1.22.0 is when Array API gains support.
+  "numpy_groupies>=0.9.16",  # Supports Python 3.7 and up.
   "pandas>=1.0.0",
+  "scikit-learn>=1.1.0",
+  "scipy>=1.7.0",
   "sortedcontainers>=2.0.0",  # Primarily used in SlidingBoundariesArchive.
-  "scikit-learn>=1.1.0",  # Primarily used in CVTArchive.
-  "scipy>=1.7.0",  # Primarily used in CVTArchive.
   "threadpoolctl>=3.0.0",
 ]
 

tests/README.md

@@ -10,15 +10,18 @@ For information on running tests, see [CONTRIBUTING.md](../CONTRIBUTING.md).
 ## External Libraries
 
 Some features of pyribs require external libraries that are optional and thus
-not specified in the default installation command. We separate these tests into
-separate directories:
-
-- `visualize_qdax/` tests visualization of QDax components
-- `emitters_pycma/` holds emitter tests that require pycma
+not specified in the default installation command. We place these tests into
+separate directories, such as `visualize_qdax/` and `emitters_pycma/`.
 
 ## Additional Tests
 
 This directory also contains:
 
 - `examples.sh`: checks that the examples work end-to-end
 - `tutorials.sh`: checks that the tutorials work end-to-end
+
+## Array API
+
+To write tests for components that feature the Array API, use the
+`xp_and_device` fixture to receive a tuple with the array namespace `xp` and the
+device `device`. `xp_and_device` is implemented in `tests/conftest.py`.

tests/conftest.py

@@ -0,0 +1,34 @@
+"""Shared functionality for all tests."""
+import numpy as np
+import pytest
+from array_api_compat import array_namespace
+
+# Array API backend handling. Adapted from scipy:
+# https://github.com/scipy/scipy/blob/888ca356eda34481e0e32b1be48c1262077d79a7/scipy/conftest.py#L283
+xp_available_backends = [
+    pytest.param(
+        (np, None),  # `None` should default to cpu for numpy.
+        id="numpy-cpu",
+    ),
+]
+
+try:
+    import torch
+
+    xp_available_backends.append(
+        pytest.param((torch, torch.device("cpu")), id="torch-cpu"))
+
+    if torch.cuda.is_available():
+        xp_available_backends.append(
+            pytest.param((torch, torch.device("cuda")), id="torch-cuda"))
+except ImportError:
+    pass
+
+
+@pytest.fixture(params=xp_available_backends)
+def xp_and_device(request):
+    """Run the test that uses this fixture on each available array API library
+    and device."""
+    xp, device = request.param
+    xp = array_namespace(xp.empty(0))
+    return xp, device