Add isdtype implementation for numpy/cupy with some initial testing

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -472,10 +472,49 @@ def vecdot(x1: ndarray, x2: ndarray, /, xp, *, axis: int = -1) -> ndarray:
     res = x1_[..., None, :] @ x2_[..., None]
     return res[..., 0, 0]
 
+# isdtype is a new function in the 2022.12 array API specification.
+
+def isdtype(
+    dtype: Dtype, kind: Union[Dtype, str, Tuple[Union[Dtype, str], ...]], xp,
+    *, _tuple=True, # Disallow nested tuples
+) -> bool:
+    """
+    Returns a boolean indicating whether a provided dtype is of a specified data type ``kind``.
+
+    See
+    https://data-apis.org/array-api/latest/API_specification/generated/array_api.isdtype.html
+    for more details
+    """
+    if isinstance(kind, tuple) and _tuple:
+        return any(isdtype(dtype, k, xp, _tuple=False) for k in kind)
+    elif isinstance(kind, str):
+        if kind == 'bool':
+            return dtype == xp.bool_
+        elif kind == 'signed integer':
+            return xp.issubdtype(dtype, xp.signedinteger)
+        elif kind == 'unsigned integer':
+            return xp.issubdtype(dtype, xp.unsignedinteger)
+        elif kind == 'integral':
+            return xp.issubdtype(dtype, xp.integer)
+        elif kind == 'real floating':
+            return xp.issubdtype(dtype, xp.floating)
+        elif kind == 'complex floating':
+            return xp.issubdtype(dtype, xp.complexfloating)
+        elif kind == 'numeric':
+            return xp.issubdtype(dtype, xp.number)
+        else:
+            raise ValueError(f"Unrecognized data type kind: {kind!r}")
+    else:
+        # This will allow things that aren't required by the spec, like
+        # isdtype(np.float64, float) or isdtype(np.int64, 'l'). Should we be
+        # more strict here to match the type annotation? Note that the
+        # numpy.array_api implementation will be very strict.
+        return dtype == kind
+
 __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'astype', 'std', 'var', 'permute_dims', 'reshape', 'argsort',
            'sort', 'sum', 'prod', 'ceil', 'floor', 'trunc', 'matmul',
-           'matrix_transpose', 'tensordot', 'vecdot']
+           'matrix_transpose', 'tensordot', 'vecdot', 'isdtype']

array_api_compat/cupy/_aliases.py

@@ -61,6 +61,7 @@
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
 vecdot = get_xp(cp)(_aliases.vecdot)
+isdtype = get_xp(cp)(_aliases.isdtype)
 
 __all__ = _aliases.__all__ + ['asarray', 'asarray_cupy', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',

array_api_compat/numpy/_aliases.py

@@ -61,6 +61,7 @@
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
 vecdot = get_xp(np)(_aliases.vecdot)
+isdtype = get_xp(np)(_aliases.isdtype)
 
 __all__ = _aliases.__all__ + ['asarray', 'asarray_numpy', 'bool', 'acos',
                               'acosh', 'asin', 'asinh', 'atan', 'atan2',

tests/_helpers.py

@@ -0,0 +1,8 @@
+from importlib import import_module
+
+import pytest
+
+def import_(library):
+    if library == 'cupy':
+        return pytest.importorskip(library)
+    return import_module(library)

tests/test_array_namespace.py

@@ -1,12 +1,15 @@
 import array_api_compat
 from array_api_compat import array_namespace
+
+from ._helpers import import_
+
 import pytest
 
 
 @pytest.mark.parametrize("library", ["cupy", "numpy", "torch"])
 @pytest.mark.parametrize("api_version", [None, '2021.12'])
 def test_array_namespace(library, api_version):
-    lib = pytest.importorskip(library)
+    lib = import_(library)
 
     array = lib.asarray([1.0, 2.0, 3.0])
     namespace = array_api_compat.array_namespace(array, api_version=api_version)

tests/test_isdtype.py

@@ -0,0 +1,88 @@
+"""
+isdtype is not yet tested in the test suite, and it should extend properly to
+non-spec dtypes
+"""
+
+from ._helpers import import_
+
+import pytest
+
+# Check the known dtypes by their string names
+
+def _spec_dtypes(library):
+    if library == 'torch':
+        # torch does not have unsigned integer dtypes
+        return {
+            'bool',
+            'uint8',
+            'int8',
+            'int16',
+            'int32',
+            'int64',
+            'float32',
+            'float64',
+        }
+    else:
+        return {
+            'bool',
+            'float32',
+            'float64',
+            'int16',
+            'int32',
+            'int64',
+            'int8',
+            'uint16',
+            'uint32',
+            'uint64',
+            'uint8',
+        }
+
+dtype_categories = {
+    'bool': lambda d: d == 'bool',
+    'signed integer': lambda d: d.startswith('int'),
+    'unsigned integer': lambda d: d.startswith('uint'),
+    'integral': lambda d: dtype_categories['signed integer'](d) or
+                          dtype_categories['unsigned integer'](d),
+    'real floating': lambda d: d.startswith('float'),
+    'complex floating': lambda d: d.startswith('complex'),
+    'numeric': lambda d: dtype_categories['integral'](d) or
+                         dtype_categories['real floating'](d) or
+                         dtype_categories['complex floating'](d),
+}
+
+def isdtype_(dtype_, kind):
+    # Check a dtype_ string against kind. Note that 'bool' technically has two
+    # meanings here but they are both the same.
+    if kind in dtype_categories:
+        res = dtype_categories[kind](dtype_)
+    else:
+        res = dtype_ == kind
+    assert type(res) is bool
+    return res
+
+@pytest.mark.parametrize("library", ["cupy", "numpy", "torch"])
+def test_isdtype_spec_dtypes(library):
+    xp = import_('array_api_compat.' + library)
+
+    isdtype = xp.isdtype
+
+    for dtype_ in _spec_dtypes(library):
+        for dtype2_ in _spec_dtypes(library):
+            dtype = getattr(xp, dtype_)
+            dtype2 = getattr(xp, dtype2_)
+            res = isdtype_(dtype_, dtype2_)
+            assert isdtype(dtype, dtype2) is res, (dtype_, dtype2_)
+
+        for cat in dtype_categories:
+            res = isdtype_(dtype_, cat)
+            assert isdtype(dtype, cat) == res, (dtype_, cat)
+
+        # Basic tuple testing (the array-api testsuite will be more complete here)
+        for kind1_ in [*_spec_dtypes(library), *dtype_categories]:
+            for kind2_ in [*_spec_dtypes(library), *dtype_categories]:
+                kind1 = kind1_ if kind1_ in dtype_categories else getattr(xp, kind1_)
+                kind2 = kind2_ if kind2_ in dtype_categories else getattr(xp, kind2_)
+                kind = (kind1, kind2)
+
+                res = isdtype_(dtype_, kind1_) or isdtype_(dtype_, kind2_)
+                assert isdtype(dtype, kind) == res, (dtype_, (kind1_, kind2_))