Make is_numpy_array, is_cupy_array, is_torch_array, and is_dask_array public

Author: asmeurer

README.md

@@ -99,6 +99,11 @@ part of the specification but which are useful for using the array API:
 - `is_array_api_obj(x)`: Return `True` if `x` is an array API compatible array
   object.
 
+- `is_numpy_array(x)`, `is_cupy_array(x)`, `is_torch_array(x)`,
+  `is_dask_array(x)`: return `True` if `x` is an array from the corresponding
+  library. These functions do not import the underlying library if it has not
+  already been imported, so they are cheap to use.
+
 - `array_namespace(*xs)`: Get the corresponding array API namespace for the
   arrays `xs`. For example, if the arrays are NumPy arrays, the returned
   namespace will be `array_api_compat.numpy`. Note that this function will

array_api_compat/common/_aliases.py

@@ -13,7 +13,7 @@
 from types import ModuleType
 import inspect
 
-from ._helpers import _check_device, _is_numpy_array, array_namespace
+from ._helpers import _check_device, is_numpy_array, array_namespace
 
 # These functions are modified from the NumPy versions.
 
@@ -309,7 +309,7 @@ def _asarray(
         raise ValueError("Unrecognized namespace argument to asarray()")
 
     _check_device(xp, device)
-    if _is_numpy_array(obj):
+    if is_numpy_array(obj):
         import numpy as np
         if hasattr(np, '_CopyMode'):
             # Not present in older NumPys

array_api_compat/common/_helpers.py

@@ -10,7 +10,7 @@
 import sys
 import math
 
-def _is_numpy_array(x):
+def is_numpy_array(x):
     # Avoid importing NumPy if it isn't already
     if 'numpy' not in sys.modules:
         return False
@@ -20,7 +20,7 @@ def _is_numpy_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, (np.ndarray, np.generic))
 
-def _is_cupy_array(x):
+def is_cupy_array(x):
     # Avoid importing NumPy if it isn't already
     if 'cupy' not in sys.modules:
         return False
@@ -30,7 +30,7 @@ def _is_cupy_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, (cp.ndarray, cp.generic))
 
-def _is_torch_array(x):
+def is_torch_array(x):
     # Avoid importing torch if it isn't already
     if 'torch' not in sys.modules:
         return False
@@ -40,7 +40,7 @@ def _is_torch_array(x):
     # TODO: Should we reject ndarray subclasses?
     return isinstance(x, torch.Tensor)
 
-def _is_dask_array(x):
+def is_dask_array(x):
     # Avoid importing dask if it isn't already
     if 'dask.array' not in sys.modules:
         return False
@@ -53,10 +53,10 @@ 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 _is_torch_array(x) \
-        or _is_dask_array(x) \
+    return is_numpy_array(x) \
+        or is_cupy_array(x) \
+        or is_torch_array(x) \
+        or is_dask_array(x) \
         or hasattr(x, '__array_namespace__')
 
 def _check_api_version(api_version):
@@ -81,31 +81,31 @@ def your_function(x, y):
     """
     namespaces = set()
     for x in xs:
-        if _is_numpy_array(x):
+        if is_numpy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import numpy as numpy_namespace
                 namespaces.add(numpy_namespace)
             else:
                 import numpy as np
                 namespaces.add(np)
-        elif _is_cupy_array(x):
+        elif is_cupy_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import cupy as cupy_namespace
                 namespaces.add(cupy_namespace)
             else:
                 import cupy as cp
                 namespaces.add(cp)
-        elif _is_torch_array(x):
+        elif is_torch_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from .. import torch as torch_namespace
                 namespaces.add(torch_namespace)
             else:
                 import torch
                 namespaces.add(torch)
-        elif _is_dask_array(x):
+        elif is_dask_array(x):
             _check_api_version(api_version)
             if _use_compat:
                 from ..dask import array as dask_namespace
@@ -156,7 +156,7 @@ def device(x: "Array", /) -> "Device":
     out: device
         a ``device`` object (see the "Device Support" section of the array API specification).
     """
-    if _is_numpy_array(x):
+    if is_numpy_array(x):
         return "cpu"
     return x.device
 
@@ -225,18 +225,18 @@ def to_device(x: "Array", device: "Device", /, *, stream: "Optional[Union[int, A
     .. 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 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):
+    elif is_cupy_array(x):
         # cupy does not yet have to_device
         return _cupy_to_device(x, device, stream=stream)
-    elif _is_torch_array(x):
+    elif is_torch_array(x):
         return _torch_to_device(x, device, stream=stream)
-    elif _is_dask_array(x):
+    elif is_dask_array(x):
         if stream is not None:
             raise ValueError("The stream argument to to_device() is not supported")
         # TODO: What if our array is on the GPU already?
@@ -253,4 +253,6 @@ def size(x):
         return None
     return math.prod(x.shape)
 
-__all__ = ['is_array_api_obj', 'array_namespace', 'get_namespace', 'device', 'to_device', 'size']
+__all__ = ['is_array_api_obj', 'array_namespace', 'get_namespace', 'device',
+           'to_device', 'size', 'is_numpy_array', 'is_cupy_array',
+           'is_torch_array', 'is_dask_array']

tests/test_helpers.py

@@ -0,0 +1,25 @@
+from array_api_compat import (is_numpy_array, is_cupy_array, is_torch_array,
+                              is_dask_array, is_array_api_obj)
+
+from ._helpers import import_
+
+import pytest
+
+is_functions = {
+    'numpy': 'is_numpy_array',
+    'cupy': 'is_cupy_array',
+    'torch': 'is_torch_array',
+    'dask.array': 'is_dask_array',
+}
+
+@pytest.mark.parametrize('library', is_functions.keys())
+@pytest.mark.parametrize('func', is_functions.values())
+def test_is_xp_array(library, func):
+    lib = import_(library)
+    is_func = globals()[func]
+
+    x = lib.asarray([1, 2, 3])
+
+    assert is_func(x) == (func == is_functions[library])
+
+    assert is_array_api_obj(x)