ENH: asarray for Array API support

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,165 @@
+"""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
+# probably want to vendor it (submodule)
+import array_api_compat
+import array_api_compat.numpy
+
+__all__ = ['array_namespace', 'asarray', 'asarray_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")
+
+
+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
+
+    compliance_scipy(*arrays)
+
+    return array_api_compat.array_namespace(*arrays)
+
+
+def asarray(array, dtype=None, order=None, copy=None, *, xp=None):
+    """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.
+        return xp.asarray(array)
+    else:
+        return xp.asarray(array, dtype=dtype, copy=copy)
+
+
+def asarray_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 = asarray_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(

scipy/_lib/tests/test_array_api.py

@@ -0,0 +1,75 @@
+import numpy as np
+from numpy.testing import assert_equal
+import pytest
+
+from scipy.conftest import array_api_compatible
+from scipy._lib._array_api import (
+    _GLOBAL_CONFIG, array_namespace, asarray, asarray_namespace,
+    to_numpy
+)
+
+
+if not _GLOBAL_CONFIG["SCIPY_ARRAY_API"]:
+    pytest.skip(
+        "Array API test; set environment variable array_api_dispatch=1 to run it",
+        allow_module_level=True
+    )
+
+
+def test_array_namespace():
+    x, y = np.array([0, 1, 2]), np.array([0, 1, 2])
+    xp = array_namespace(x, y)
+    assert xp.__name__ == 'array_api_compat.numpy'
+
+    _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = False
+    xp = array_namespace(x, y)
+    assert xp.__name__ == 'numpy'
+    _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = True
+
+
+@array_api_compatible
+def test_asarray(xp):
+    x, y = asarray([0, 1, 2], xp=xp), asarray(np.arange(3), xp=xp)
+    ref = np.array([0, 1, 2])
+    assert_equal(x, ref)
+    assert_equal(y, ref)
+
+
+def test_asarray_namespace():
+    x, y = np.array([0, 1, 2]), np.array([0, 1, 2])
+    x, y, xp_ = asarray_namespace(x, y)
+    assert xp_.__name__ == 'array_api_compat.numpy'
+    ref = np.array([0, 1, 2])
+    assert_equal(x, ref)
+    assert_equal(y, ref)
+    assert type(x) == type(y)
+
+    _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = False
+    x, y, xp_ = asarray_namespace(x, y)
+    assert xp_.__name__ == 'numpy'
+    _GLOBAL_CONFIG["SCIPY_ARRAY_API"] = True
+
+
+@array_api_compatible
+def test_to_numpy(xp):
+    x = xp.asarray([0, 1, 2])
+    x = to_numpy(x, xp=xp)
+    assert isinstance(x, np.ndarray)
+
+
+@pytest.mark.filterwarnings("ignore: the matrix subclass")
+def test_raises():
+    msg = "'numpy.ma.MaskedArray' are not supported"
+    with pytest.raises(TypeError, match=msg):
+        array_namespace(np.ma.array(1), np.array(1))
+
+    msg = "'numpy.matrix' are not supported"
+    with pytest.raises(TypeError, match=msg):
+        array_namespace(np.array(1), np.matrix(1))
+
+    msg = "Only support Array API"
+    with pytest.raises(TypeError, match=msg):
+        array_namespace([0, 1, 2])
+
+    with pytest.raises(TypeError, match=msg):
+        array_namespace(1)

scipy/conftest.py

@@ -4,10 +4,13 @@
 import warnings
 
 import numpy as np
+import numpy.array_api
 import numpy.testing as npt
+
 from scipy._lib._fpumode import get_fpu_mode
 from scipy._lib._testutils import FPUModeChangeWarning
 from scipy._lib import _pep440
+from scipy._lib._array_api import SCIPY_ARRAY_API
 
 
 def pytest_configure(config):
@@ -93,3 +96,7 @@ def check_fpu_mode(request):
         warnings.warn("FPU mode changed from {:#x} to {:#x} during "
                       "the test".format(old_mode, new_mode),
                       category=FPUModeChangeWarning, stacklevel=0)
+
+array_api_compatible = pytest.mark.parametrize(
+    "xp", [np, *((numpy.array_api,) if SCIPY_ARRAY_API else ())]
+)