ENH: add machinery to support Array API (scipy#18668)

See api-dev doc updates in this commit for details.

Co-authored-by: Ralf Gommers <[email protected]>
Co-authored-by: Tyler Reddy <[email protected]>

Author: 3 people

.github/workflows/array_api.yml

@@ -0,0 +1,89 @@
+name: Array API
+
+on:
+  push:
+    branches:
+      - maintenance/**
+  pull_request:
+    branches:
+      - main
+      - maintenance/**
+
+permissions:
+   contents: read  # to fetch code (actions/checkout)
+
+env:
+  CCACHE_DIR: "${{ github.workspace }}/.ccache"
+  INSTALLDIR: "build-install"
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  pytorch_cpu:
+    name: Linux PyTorch CPU
+    if: "github.repository == 'scipy/scipy' || github.repository == ''"
+    runs-on: ubuntu-22.04
+    strategy:
+      matrix:
+        python-version: ['3.11']
+        maintenance-branch:
+          - ${{ contains(github.ref, 'maintenance/') || contains(github.base_ref, 'maintenance/') }}
+        exclude:
+          - maintenance-branch: true
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        submodules: recursive
+
+    - name: Setup Python
+      uses: actions/setup-python@v4
+      with:
+        python-version: ${{ matrix.python-version }}
+        cache: 'pip'  # not using a path to also cache pytorch
+
+    - name: Install Ubuntu dependencies
+      run: |
+        sudo apt-get update
+        sudo apt-get install -y libopenblas-dev libatlas-base-dev liblapack-dev gfortran libgmp-dev libmpfr-dev libsuitesparse-dev ccache libmpc-dev
+
+    - name: Install Python packages
+      run: |
+        python -m pip install numpy cython pytest pytest-xdist pytest-timeout pybind11 mpmath gmpy2 pythran ninja meson click rich-click doit pydevtool pooch
+
+    - name: Install PyTorch CPU
+      run: |
+        python -m pip install "torch<2.1" --index-url https://download.pytorch.org/whl/cpu
+
+    - name:  Prepare compiler cache
+      id:    prep-ccache
+      shell: bash
+      run: |
+        mkdir -p "${CCACHE_DIR}"
+        echo "dir=$CCACHE_DIR" >> $GITHUB_OUTPUT
+        NOW=$(date -u +"%F-%T")
+        echo "timestamp=${NOW}" >> $GITHUB_OUTPUT
+
+    - name: Setup compiler cache
+      uses:  actions/cache@v3
+      id:    cache-ccache
+      with:
+        path: ${{ steps.prep-ccache.outputs.dir }}
+        key:  ${{ github.workflow }}-${{ matrix.python-version }}-ccache-linux-${{ steps.prep-ccache.outputs.timestamp }}
+        restore-keys: |
+          ${{ github.workflow }}-${{ matrix.python-version }}-ccache-linux-
+
+    - name: Setup build and install scipy
+      run: |
+        python dev.py build
+
+    - name: Test SciPy
+      run: |
+        export OMP_NUM_THREADS=2
+        export SCIPY_USE_PROPACK=1
+        # expand as new modules are added
+        # --array-api-backend == -b
+        python dev.py --no-build test -b pytorch -b numpy -b numpy.array_api -s cluster -- --durations 10 --timeout=60
+        python dev.py --no-build test --array-api-backend all --tests scipy/_lib/tests/test_array_api.py

.gitmodules

@@ -17,3 +17,6 @@
 	path = scipy/_lib/boost_math
 	url = https://github.com/boostorg/math.git
 	shallow = true
+[submodule "scipy/_lib/array_api_compat"]
+	path = scipy/_lib/array_api_compat
+	url = https://github.com/data-apis/array-api-compat.git

LICENSES_bundled.txt

@@ -256,3 +256,8 @@ Files: scipy/stats/_rcont/[logfactorial.h,logfactorial.c]
 License 3-Clause BSD
   For details, see header inside scipy/stats/_rcont/logfactorial.h
   and scipy/stats/_rcont/logfactorial.c
+
+Name: array-api-compat
+Files: scipy/_lib/array-api-compat/*
+License: MIT
+  For details, see scipy/_lib/array-api-compat/LICENCE

benchmarks/benchmarks/cluster.py

@@ -50,7 +50,7 @@ def time_kmeans2(self, k, init):
 
 
 class VQ(Benchmark):
-    params = [[2, 10, 50], ['float32', 'float64', 'float128']]
+    params = [[2, 10, 50], ['float32', 'float64']]
     param_names = ['k', 'dtype']
 
     def __init__(self):

dev.py

@@ -685,6 +685,7 @@ class Test(Task):
     $ python dev.py test -t scipy.optimize.tests.test_minimize_constrained
     $ python dev.py test -s cluster -m full --durations 20
     $ python dev.py test -s stats -- --tb=line  # `--` passes next args to pytest
+    $ python dev.py test -b numpy -b pytorch -s cluster
     ```
     """  # noqa: E501
     ctx = CONTEXT
@@ -716,6 +717,13 @@ class Test(Task):
         ['--parallel', '-j'], default=1, metavar='N_JOBS',
         help="Number of parallel jobs for testing"
     )
+    array_api_backend = Option(
+        ['--array-api-backend', '-b'], default=None, metavar='ARRAY_BACKEND',
+        multiple=True,
+        help=(
+            "Array API backend ('all', 'numpy', 'pytorch', 'cupy', 'numpy.array_api')."
+        )
+    )
     # Argument can't have `help=`; used to consume all of `-- arg1 arg2 arg3`
     pytest_args = Argument(
         ['pytest_args'], nargs=-1, metavar='PYTEST-ARGS', required=False
@@ -756,6 +764,9 @@ def scipy_tests(cls, args, pytest_args):
         else:
             tests = None
 
+        if len(args.array_api_backend) != 0:
+            os.environ['SCIPY_ARRAY_API'] = json.dumps(list(args.array_api_backend))
+
         runner, version, mod_path = get_test_runner(PROJECT_MODULE)
         # FIXME: changing CWD is not a good practice
         with working_dir(dirs.site):

doc/source/dev/api-dev/api-dev-toc.rst

@@ -9,3 +9,4 @@ SciPy API Development Guide
 
    nan_policy
    special_ufuncs
+   array_api