Add a cupy submodule

Full support still needs to be tested, and also we need to double check if any
of the aliases are unnecessary for cupy.

Author: asmeurer

numpy_array_api_compat/common/_aliases.py

@@ -12,7 +12,7 @@
 from typing import NamedTuple
 from types import ModuleType
 
-from ..numpy._helpers import _is_numpy_array, get_namespace
+from ._helpers import _check_device, _is_numpy_array, get_namespace
 from .._internal import get_xp
 
 # Basic renames
@@ -189,10 +189,6 @@ def permute_dims(x: ndarray, /, xp, axes: Tuple[int, ...]) -> ndarray:
 
 # Creation functions add the device keyword (which does nothing for NumPy)
 
-def _check_device(device):
-    if device not in ["cpu", None]:
-        raise ValueError(f"Unsupported device {device!r}")
-
 # asarray also adds the copy keyword
 def _asarray(
     obj: Union[
@@ -232,7 +228,7 @@ def _asarray(
     else:
         raise ValueError("Unrecognized namespace argument to asarray()")
 
-    _check_device(device)
+    _check_device(xp, device)
     if _is_numpy_array(obj):
         import numpy as np
         COPY_FALSE = (False, np._CopyMode.IF_NEEDED)
@@ -263,7 +259,7 @@ def arange(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.arange(start, stop=stop, step=step, dtype=dtype)
 
 @get_xp
@@ -274,14 +270,14 @@ def empty(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.empty(shape, dtype=dtype)
 
 @get_xp
 def empty_like(
     x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.empty_like(x, dtype=dtype)
 
 @get_xp
@@ -295,7 +291,7 @@ def eye(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.eye(n_rows, M=n_cols, k=k, dtype=dtype)
 
 @get_xp
@@ -307,7 +303,7 @@ def full(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.full(shape, fill_value, dtype=dtype)
 
 @get_xp
@@ -320,7 +316,7 @@ def full_like(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.full_like(x, fill_value, dtype=dtype)
 
 @get_xp
@@ -335,7 +331,7 @@ def linspace(
     device: Optional[Device] = None,
     endpoint: bool = True,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.linspace(start, stop, num, dtype=dtype, endpoint=endpoint)
 
 @get_xp
@@ -346,14 +342,14 @@ def ones(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.ones(shape, dtype=dtype)
 
 @get_xp
 def ones_like(
     x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.ones_like(x, dtype=dtype)
 
 @get_xp
@@ -364,14 +360,14 @@ def zeros(
     dtype: Optional[Dtype] = None,
     device: Optional[Device] = None,
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.zeros(shape, dtype=dtype)
 
 @get_xp
 def zeros_like(
     x: ndarray, /, xp, *, dtype: Optional[Dtype] = None, device: Optional[Device] = None
 ) -> ndarray:
-    _check_device(device)
+    _check_device(xp, device)
     return xp.zeros_like(x, dtype=dtype)
 
 # xp.reshape calls the keyword argument 'newshape' instead of 'shape'

numpy_array_api_compat/common/_helpers.py

@@ -1,6 +1,40 @@
 """
 Various helper functions which are not part of the spec.
+
+Functions which start with an underscore are for internal use only but helpers
+that are in __all__ are intended as additional helper functions for use by end
+users of the compat library.
 """
+from __future__ import annotations
+
+import sys
+
+def _is_numpy_array(x):
+    # Avoid importing NumPy if it isn't already
+    if 'numpy' not in sys.modules:
+        return False
+
+    import numpy as np
+
+    # TODO: Should we reject ndarray subclasses?
+    return isinstance(x, (np.ndarray, np.generic))
+
+def _is_cupy_array(x):
+    # Avoid importing NumPy if it isn't already
+    if 'cupy' not in sys.modules:
+        return False
+
+    import cupy as cp
+
+    # TODO: Should we reject ndarray subclasses?
+    return isinstance(x, (cp.ndarray, cp.generic))
+
+def is_array_api_obj(x):
+    """
+    Check if x is an array API compatible array object.
+    """
+    return _is_numpy_array(x) or _is_cupy_array(x) or hasattr(x, '__array_namespace__')
+
 def get_namespace(*xs, _use_compat=True):
     """
     Get the array API compatible namespace for the arrays `xs`.
@@ -35,3 +69,102 @@ def get_namespace(*xs, _use_compat=True):
     xp, = namespaces
 
     return xp
+
+
+def _check_device(xp, device):
+    if xp == sys.modules.get('numpy'):
+        if device not in ["cpu", None]:
+            raise ValueError(f"Unsupported device for NumPy: {device!r}")
+
+# device() is not on numpy.ndarray and and to_device() is not on numpy.ndarray
+# or cupy.ndarray. They are not included in array objects of this library
+# because this library just reuses the respective ndarray classes without
+# wrapping or subclassing them. These helper functions can be used instead of
+# the wrapper functions for libraries that need to support both NumPy/CuPy and
+# other libraries that use devices.
+def device(x: "Array", /) -> "Device":
+    """
+    Hardware device the array data resides on.
+
+    Parameters
+    ----------
+    x: array
+        array instance from NumPy or an array API compatible library.
+
+    Returns
+    -------
+    out: device
+        a ``device`` object (see the "Device Support" section of the array API specification).
+    """
+    if _is_numpy_array(x):
+        return "cpu"
+    return x.device
+
+# Based on cupy.array_api.Array.to_device
+def _cupy_to_device(x, device, /, stream=None):
+    import cupy as cp
+    from cupy.cuda import Device as _Device
+    from cupy.cuda import stream as stream_module
+    from cupy_backends.cuda.api import runtime
+
+    if device == x.device:
+        return x
+    elif not isinstance(device, _Device):
+        raise ValueError(f"Unsupported device {device!r}")
+    else:
+        # see cupy/cupy#5985 for the reason how we handle device/stream here
+        prev_device = runtime.getDevice()
+        prev_stream: stream_module.Stream = None
+        if stream is not None:
+            prev_stream = stream_module.get_current_stream()
+            # stream can be an int as specified in __dlpack__, or a CuPy stream
+            if isinstance(stream, int):
+                stream = cp.cuda.ExternalStream(stream)
+            elif isinstance(stream, cp.cuda.Stream):
+                pass
+            else:
+                raise ValueError('the input stream is not recognized')
+            stream.use()
+        try:
+            runtime.setDevice(device.id)
+            arr = x.copy()
+        finally:
+            runtime.setDevice(prev_device)
+            if stream is not None:
+                prev_stream.use()
+        return arr
+
+def to_device(x: "Array", device: "Device", /, *, stream: Optional[Union[int, Any]] = None) -> "Array":
+    """
+    Copy the array from the device on which it currently resides to the specified ``device``.
+
+    Parameters
+    ----------
+    x: array
+        array instance from NumPy or an array API compatible library.
+    device: device
+        a ``device`` object (see the "Device Support" section of the array API specification).
+    stream: Optional[Union[int, Any]]
+        stream object to use during copy. In addition to the types supported in ``array.__dlpack__``, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable.
+
+    Returns
+    -------
+    out: array
+        an array with the same data and data type as ``x`` and located on the specified ``device``.
+
+    .. note::
+       If ``stream`` is given, the copy operation should be enqueued on the provided ``stream``; otherwise, the copy operation should be enqueued on the default stream/queue. Whether the copy is performed synchronously or asynchronously is implementation-dependent. Accordingly, if synchronization is required to guarantee data safety, this must be clearly explained in a conforming library's documentation.
+    """
+    if _is_numpy_array(x):
+        if stream is not None:
+            raise ValueError("The stream argument to to_device() is not supported")
+        if device == 'cpu':
+            return x
+        raise ValueError(f"Unsupported device {device!r}")
+    elif _is_cupy_array(x):
+        # cupy does not yet have to_device
+        return _cupy_to_device(x, device, stream=stream)
+
+    return x.to_device(device, stream=stream)
+
+__all__ = ['is_array_api_obj', 'get_namespace', 'device', 'to_device']

numpy_array_api_compat/cupy/__init__.py

@@ -11,10 +11,10 @@
 #
 # from . import linalg
 #
-# It doesn't overwrite np.linalg from above. The import is generated
+# It doesn't overwrite cupy.linalg from above. The import is generated
 # dynamically so that the library can be vendored.
 __import__(__package__ + '.linalg')
 
 from .linalg import matrix_transpose, vecdot
 
-from ._helpers import *
+from ..common._helpers import *

numpy_array_api_compat/cupy/_aliases.py

@@ -0,0 +1,11 @@
+from functools import partial
+
+from ..common._aliases import *
+from ..common._aliases import _asarray
+from ..common._aliases import __all__
+
+asarray = asarray_cupy = partial(_asarray, namespace='cupy')
+asarray.__doc__ = _asarray.__doc__
+del partial
+
+__all__ = __all__ + ['asarray', 'asarray_cupy']

numpy_array_api_compat/cupy/linalg.py

@@ -0,0 +1,14 @@
+from cupy.linalg import *
+# cupy.linalg doesn't have __all__. If it is added, replace this with
+#
+# from cupy.linalg import __all__ as linalg_all
+_n = {}
+exec('from cupy.linalg import *', _n)
+del _n['__builtins__']
+linalg_all = list(_n)
+del _n
+
+from ..common.linalg import *
+from ..common.linalg import __all__ as common_linalg_all
+
+__all__ = linalg_all + common_linalg_all

numpy_array_api_compat/numpy/__init__.py

@@ -17,4 +17,4 @@
 
 from .linalg import matrix_transpose, vecdot
 
-from ._helpers import *
+from ..common._helpers import *