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 *