Add device() and to_device() helper functions

Author: asmeurer

numpy_array_api_compat/__init__.py

@@ -16,18 +16,26 @@
 to ensure they are not using functionality outside of the standard, but prefer
 this implementation for the default when working with NumPy arrays.
 
+In addition, several helper functions are provided in this library which are
+not part of the array API specification but which are useful for libraries
+writing against the array API specification who wish to support NumPy and
+other array API compatible libraries.
+
 Known differences from the Array API spec:
 
 - The array methods __array_namespace__, device, to_device, and mT are not
   defined. This reuses np.ndarray and we don't want to monkeypatch or wrap it.
+  The helper functions device() and to_device() are provided to work around
+  these missing methods. x.mT can be replaced with
+  xp.linalg.matrix_transpose(x).
 
 - NumPy value-based casting for scalars will be in effect unless explicitly
   disabled with the environment variable NPY_PROMOTION_STATE=weak or
   np._set_promotion_state('weak') (requires NumPy 1.24 or newer, see NEP 50
   and https://github.com/numpy/numpy/issues/22341)
 
-- NumPy functions which are not wrapped may not have the same type aliases as
-  the spec.
+- NumPy functions which are not wrapped may not have the same type annotations
+  as the spec.
 
 - NumPy functions which are not wrapped may not use positional-only arguments.
 
@@ -46,3 +54,5 @@
 import numpy_array_api_compat.linalg
 
 from .linalg import matrix_transpose, vecdot
+
+from .helpers import *

numpy_array_api_compat/_helpers.py

@@ -0,0 +1,61 @@
+"""
+Various helper functions which are not part of the spec.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+
+# device and to_device are not included in array object of this library
+# because this library just reuses ndarray without wrapping or subclassing it.
+# These helper functions can be used instead of the wrapper functions for
+# libraries that need to support both NumPy 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 isinstance(x, np.ndarray):
+        return "cpu"
+    return x.device
+
+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 isinstance(x, np.ndarray):
+        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}")
+
+    return x.to_device(device, stream=stream)
+
+__all__ = ['device', 'to_device']