ENH: asarray for Array API support

.github/workflows/array_api.yml

@@ -0,0 +1,85 @@
+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
+        # Packages for Array API testing
+        python -m pip install array-api-compat
+        python -m pip install torch --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
+        python dev.py --no-build test --array-api-backend pytorch -s cluster -- --durations 10 --timeout=60

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', '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):

pytest.ini

@@ -14,3 +14,4 @@ filterwarnings =
     ignore:.*The distutils.* is deprecated.*:DeprecationWarning
     ignore:\s*.*numpy.distutils.*:DeprecationWarning
     ignore:.*The --rsyncdir command line argument.*:DeprecationWarning
+    ignore:.*The numpy.array_api submodule is still experimental.*:UserWarning

scipy/_lib/_array_api.py

@@ -0,0 +1,187 @@
+"""Utility functions to use Python Array API compatible libraries.
+
+For the context about the Array API see:
+https://data-apis.org/array-api/latest/purpose_and_scope.html
+
+The SciPy use case of the Array API is described on the following page:
+https://data-apis.org/array-api/latest/use_cases.html#use-case-scipy
+"""
+import os
+
+import numpy as np
+from numpy.core.numerictypes import typecodes
+# probably want to vendor it (submodule)
+import array_api_compat
+import array_api_compat.numpy
+
+__all__ = ['array_namespace', 'as_xparray', 'as_xparray_namespace']
+
+
+# SCIPY_ARRAY_API, array_api_dispatch is used by sklearn
+array_api_dispatch = os.environ.get("array_api_dispatch", False)
+SCIPY_ARRAY_API = os.environ.get("SCIPY_ARRAY_API", array_api_dispatch)
+
+_GLOBAL_CONFIG = {"SCIPY_ARRAY_API": SCIPY_ARRAY_API}
+
+
+def compliance_scipy(*arrays):
+    """Raise exceptions on known-bad subclasses.
+
+    The following subclasses are not supported and raise and error:
+    - `np.ma.MaskedArray`
+    - `numpy.matrix`
+    - Any array-like which is not Array API compatible
+    """
+    for array in arrays:
+        if isinstance(array, np.ma.MaskedArray):
+            raise TypeError("'numpy.ma.MaskedArray' are not supported")
+        elif isinstance(array, np.matrix):
+            raise TypeError("'numpy.matrix' are not supported")
+        elif not array_api_compat.is_array_api_obj(array):
+            raise TypeError("Only support Array API compatible arrays")
+        elif array.dtype is np.dtype('O'):
+            raise ValueError('object arrays are not supported')
+
+
+def _check_finite(array, xp):
+    """Check for NaNs or Infs."""
+    msg = "array must not contain infs or NaNs"
+    try:
+        if not xp.isfinite(array).all():
+            raise ValueError(msg)
+    except TypeError:
+        raise ValueError(msg)
+
+
+def array_namespace(*arrays):
+    """Get the array API compatible namespace for the arrays xs.
+
+    Parameters
+    ----------
+    *arrays : sequence of array_like
+        Arrays used to infer the common namespace.
+
+    Returns
+    -------
+    namespace : module
+        Common namespace.
+
+    Notes
+    -----
+    Thin wrapper around `array_api_compat.array_namespace`.
+
+    1. Check for the global switch: SCIPY_ARRAY_API. This can also be accessed
+       dynamically through ``_GLOBAL_CONFIG['SCIPY_ARRAY_API']``.
+    2. `compliance_scipy` raise exceptions on known-bad subclasses. See
+       it's definition for more details.
+
+    When the global switch is False, it defaults to the `numpy` namespace.
+    In that case, there is no compliance check. This is a convenience to
+    ease the adoption. Otherwise, arrays must comply with the new rules.
+    """
+    if not _GLOBAL_CONFIG["SCIPY_ARRAY_API"]:
+        # here we could wrap the namespace if needed
+        return np
+
+    arrays = [array for array in arrays if array is not None]
+
+    compliance_scipy(*arrays)
+
+    return array_api_compat.array_namespace(*arrays)
+
+
+def as_xparray(
+    array, dtype=None, order=None, copy=None, *, xp=None, check_finite=True
+):
+    """Drop-in replacement for `np.asarray`.
+
+    Memory layout parameter `order` is not exposed in the Array API standard.
+    `order` is only enforced if the input array implementation
+    is NumPy based, otherwise `order` is just silently ignored.
+
+    The purpose of this helper is to make it possible to share code for data
+    container validation without memory copies for both downstream use cases.
+    """
+    if xp is None:
+        xp = array_namespace(array)
+    if xp.__name__ in {"numpy", "array_api_compat.numpy", "numpy.array_api"}:
+        # Use NumPy API to support order
+        if copy is True:
+            array = np.array(array, order=order, dtype=dtype)
+        else:
+            array = np.asarray(array, order=order, dtype=dtype)
+
+        # At this point array is a NumPy ndarray. We convert it to an array
+        # container that is consistent with the input's namespace.
+        array = xp.asarray(array)
+    else:
+        array = xp.asarray(array, dtype=dtype, copy=copy)
+
+    if check_finite:
+        _check_finite(array, xp)
+
+    return array
+
+
+def as_xparray_namespace(*arrays):
+    """Validate and convert arrays to a common namespace.
+
+    Parameters
+    ----------
+    *arrays : sequence of array_like
+        Arrays to validate and convert.
+
+    Returns
+    -------
+    *arrays : sequence of array_like
+        Validated and converted arrays to the common namespace.
+    namespace : module
+        Common namespace.
+
+    Notes
+    -----
+    This function is meant to be called from each public function in a SciPy
+    submodule it does the following:
+
+    1. Check for the global switch: SCIPY_ARRAY_API. This can also be accessed
+       dynamically through ``_GLOBAL_CONFIG['SCIPY_ARRAY_API']``.
+    2. `compliance_scipy` raise exceptions on known-bad subclasses. See
+       it's definition for more details.
+    3. Determine the namespace, without doing any coercion of array(-like)
+       inputs.
+    4. Call `xp.asarray` on all array.
+
+    Examples
+    --------
+    >>> import numpy as np
+    >>> x, y, xp = as_xparray_namespace(np.array([0, 1, 2]), np.array([0, 1, 2]))
+    >>> xp.__name__
+    'array_api_compat.numpy'
+    >>> x, y
+    (array([0, 1, 2]), array([0, 1, 2]))
+
+    """
+    arrays = list(arrays)
+    xp = array_namespace(*arrays)
+
+    for i, array in enumerate(arrays):
+        arrays[i] = xp.asarray(array)
+
+    return *arrays, xp
+
+
+def to_numpy(array, xp):
+    """Convert `array` into a NumPy ndarray on the CPU.
+
+    ONLY FOR TESTING
+    """
+    xp_name = xp.__name__
+
+    if xp_name in {"array_api_compat.torch", "torch"}:
+        return array.cpu().numpy()
+    elif xp_name == "cupy.array_api":
+        return array._array.get()
+    elif xp_name in {"array_api_compat.cupy", "cupy"}:  # pragma: nocover
+        return array.get()
+
+    return np.asarray(array)

scipy/_lib/meson.build

@@ -98,6 +98,7 @@ py3.extension_module('messagestream',
 
 python_sources = [
   '__init__.py',
+  '_array_api.py',
   '_bunch.py',
   '_ccallback.py',
   '_disjoint_set.py',

scipy/_lib/tests/meson.build

@@ -12,6 +12,7 @@ python_sources = [
   'test_public_api.py',
   'test_tmpdirs.py',
   'test_warnings.py',
+  'test_array_api.py',
 ]
 
 py3.install_sources(