Add array API inspection utilities to dpctl.tensor (#1469)

ndgrigorian · web-flow · commit f2af75376928 · 2023-11-08T08:38:05.000-08:00
* Adds __array_namespace_info__ inspection utility

This inspection utility is coming to the array API specification in the near future

* Set __array_api_version__ to "2022.12"

* Remove --ci from array API conformity workflow

* Adds __array_namespace_info__ docstrings

Disallows dtypes for `kind` kwarg in __array_namespace_info__().dtypes

Removes  `float16` from dtypes listed by __array_namespace_info__ as per spec

Permits dpctl.tensor.Device objects in device keyword arguments in array API inspection utilities

* Adds tests for array API inspection
diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml
@@ -666,7 +666,7 @@ jobs:
           python -c "import dpctl; dpctl.lsplatform()"
           export ARRAY_API_TESTS_MODULE=dpctl.tensor
           cd /home/runner/work/array-api-tests
-          pytest --ci --json-report --json-report-file=$FILE array_api_tests/ || true
+          pytest --json-report --json-report-file=$FILE array_api_tests/ || true
       - name: Set Github environment variables
         shell: bash -l {0}
         run: |
diff --git a/dpctl/tensor/__init__.py b/dpctl/tensor/__init__.py
@@ -93,6 +93,7 @@
 from dpctl.tensor._usmarray import usm_ndarray
 from dpctl.tensor._utility_functions import all, any
 
+from ._array_api import __array_api_version__, __array_namespace_info__
 from ._clip import clip
 from ._constants import e, inf, nan, newaxis, pi
 from ._elementwise_funcs import (
@@ -335,4 +336,6 @@
     "clip",
     "logsumexp",
     "reduce_hypot",
+    "__array_api_version__",
+    "__array_namespace_info__",
 ]
diff --git a/dpctl/tensor/_array_api.py b/dpctl/tensor/_array_api.py
@@ -0,0 +1,207 @@
+#                       Data Parallel Control (dpctl)
+#
+#  Copyright 2020-2023 Intel Corporation
+#
+#  Licensed under the Apache License, Version 2.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+
+import dpctl
+import dpctl.tensor as dpt
+from dpctl.tensor._tensor_impl import (
+    default_device_complex_type,
+    default_device_fp_type,
+    default_device_index_type,
+    default_device_int_type,
+)
+
+
+def _isdtype_impl(dtype, kind):
+    if isinstance(kind, str):
+        if kind == "bool":
+            return dtype.kind == "b"
+        elif kind == "signed integer":
+            return dtype.kind == "i"
+        elif kind == "unsigned integer":
+            return dtype.kind == "u"
+        elif kind == "integral":
+            return dtype.kind in "iu"
+        elif kind == "real floating":
+            return dtype.kind == "f"
+        elif kind == "complex floating":
+            return dtype.kind == "c"
+        elif kind == "numeric":
+            return dtype.kind in "iufc"
+        else:
+            raise ValueError(f"Unrecognized data type kind: {kind}")
+
+    elif isinstance(kind, tuple):
+        return any(_isdtype_impl(dtype, k) for k in kind)
+    else:
+        raise TypeError(f"Unsupported data type kind: {kind}")
+
+
+__array_api_version__ = "2022.12"
+
+
+class Info:
+    """
+    namespace returned by `__array_namespace_info__()`
+    """
+
+    def __init__(self):
+        self._capabilities = {
+            "boolean_indexing": True,
+            "data_dependent_shapes": True,
+        }
+        self._all_dtypes = {
+            "bool": dpt.bool,
+            "float32": dpt.float32,
+            "float64": dpt.float64,
+            "complex64": dpt.complex64,
+            "complex128": dpt.complex128,
+            "int8": dpt.int8,
+            "int16": dpt.int16,
+            "int32": dpt.int32,
+            "int64": dpt.int64,
+            "uint8": dpt.uint8,
+            "uint16": dpt.uint16,
+            "uint32": dpt.uint32,
+            "uint64": dpt.uint64,
+        }
+
+    def capabilities(self):
+        """
+        Returns a dictionary of `dpctl`'s capabilities.
+
+        Returns:
+            dict:
+                dictionary of `dpctl`'s capabilities
+                - `boolean_indexing`: bool
+                - `data_dependent_shapes`: bool
+        """
+        return self._capabilities.copy()
+
+    def default_device(self):
+        """
+        Returns the default SYCL device.
+        """
+        return dpctl.select_default_device()
+
+    def default_dtypes(self, device=None):
+        """
+        Returns a dictionary of default data types for `device`.
+
+        Args:
+            device (Optional[dpctl.SyclDevice, dpctl.SyclQueue,
+            dpctl.tensor.Device]):
+                array API concept of device used in getting default data types.
+                `device` can be `None` (in which case the default device is
+                used), an instance of :class:`dpctl.SyclDevice` corresponding
+                to a non-partitioned SYCL device, an instance of
+                :class:`dpctl.SyclQueue`, or a `Device` object returned by
+                :attr:`dpctl.tensor.usm_array.device`. Default: `None`.
+
+        Returns:
+            dict:
+                a dictionary of default data types for `device`
+                    - `real floating`: dtype
+                    - `complex floating`: dtype
+                    - `integral`: dtype
+                    - `indexing`: dtype
+        """
+        if device is None:
+            device = dpctl.select_default_device()
+        elif isinstance(device, dpt.Device):
+            device = device.sycl_device
+        return {
+            "real floating": dpt.dtype(default_device_fp_type(device)),
+            "complex floating": dpt.dtype(default_device_complex_type(device)),
+            "integral": dpt.dtype(default_device_int_type(device)),
+            "indexing": dpt.dtype(default_device_index_type(device)),
+        }
+
+    def dtypes(self, device=None, kind=None):
+        """
+        Returns a dictionary of all Array API data types of a specified `kind`
+        supported by `device`
+
+        This dictionary only includes data types supported by the array API.
+
+        See [array API](array_api).
+
+        [array_api]: https://data-apis.org/array-api/latest/
+
+        Args:
+            device (Optional[dpctl.SyclDevice, dpctl.SyclQueue,
+            dpctl.tensor.Device, str]):
+                array API concept of device used in getting default data types.
+                `device` can be `None` (in which case the default device is
+                used), an instance of :class:`dpctl.SyclDevice` corresponding
+                to a non-partitioned SYCL device, an instance of
+                :class:`dpctl.SyclQueue`, or a `Device` object returned by
+                :attr:`dpctl.tensor.usm_array.device`. Default: `None`.
+
+            kind (Optional[str, Tuple[str, ...]]):
+                data type kind.
+                - if `kind` is `None`, returns a dictionary of all data types
+                supported by `device`
+                - if `kind` is a string, returns a dictionary containing the
+                data types belonging to the data type kind specified.
+                Supports:
+                    - "bool"
+                    - "signed integer"
+                    - "unsigned integer"
+                    - "integral"
+                    - "real floating"
+                    - "complex floating"
+                    - "numeric"
+                - if `kind` is a tuple, the tuple represents a union of `kind`
+                strings, and returns a dictionary containing data types
+                corresponding to the-specified union.
+                Default: `None`.
+
+        Returns:
+            dict:
+                a dictionary of the supported data types of the specified `kind`
+        """
+        if device is None:
+            device = dpctl.select_default_device()
+        elif isinstance(device, dpt.Device):
+            device = device.sycl_device
+        _fp64 = device.has_aspect_fp64
+        if kind is None:
+            return {
+                key: val
+                for key, val in self._all_dtypes.items()
+                if (key != "float64" or _fp64)
+            }
+        else:
+            return {
+                key: val
+                for key, val in self._all_dtypes.items()
+                if (key != "float64" or _fp64) and _isdtype_impl(val, kind)
+            }
+
+    def devices(self):
+        """
+        Returns a list of supported devices.
+        """
+        return dpctl.get_devices()
+
+
+def __array_namespace_info__():
+    """__array_namespace_info__()
+
+    Returns a namespace with Array API namespace inspection utilities.
+
+    """
+    return Info()
diff --git a/dpctl/tests/test_tensor_array_api_inspection.py b/dpctl/tests/test_tensor_array_api_inspection.py
@@ -0,0 +1,163 @@
+#                      Data Parallel Control (dpctl)
+#
+# Copyright 2020-2023 Intel Corporation
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import pytest
+
+import dpctl
+import dpctl.tensor as dpt
+from dpctl.tensor._tensor_impl import (
+    default_device_complex_type,
+    default_device_fp_type,
+    default_device_index_type,
+    default_device_int_type,
+)
+
+_dtypes_no_fp16_fp64 = {
+    "bool": dpt.bool,
+    "float32": dpt.float32,
+    "complex64": dpt.complex64,
+    "complex128": dpt.complex128,
+    "int8": dpt.int8,
+    "int16": dpt.int16,
+    "int32": dpt.int32,
+    "int64": dpt.int64,
+    "uint8": dpt.uint8,
+    "uint16": dpt.uint16,
+    "uint32": dpt.uint32,
+    "uint64": dpt.uint64,
+}
+
+
+class MockDevice:
+    def __init__(self, fp16: bool, fp64: bool):
+        self.has_aspect_fp16 = fp16
+        self.has_aspect_fp64 = fp64
+
+
+def test_array_api_inspection_methods():
+    info = dpt.__array_namespace_info__()
+    assert info.capabilities()
+    assert info.default_device()
+    assert info.default_dtypes()
+    assert info.devices()
+    assert info.dtypes()
+
+
+def test_array_api_inspection_default_device():
+    assert (
+        dpt.__array_namespace_info__().default_device()
+        == dpctl.select_default_device()
+    )
+
+
+def test_array_api_inspection_devices():
+    devices1 = dpt.__array_namespace_info__().devices()
+    devices2 = dpctl.get_devices()
+    assert len(devices1) == len(devices2)
+    assert devices1 == devices2
+
+
+def test_array_api_inspection_capabilities():
+    capabilities = dpt.__array_namespace_info__().capabilities()
+    assert capabilities["boolean_indexing"]
+    assert capabilities["data_dependent_shapes"]
+
+
+def test_array_api_inspection_default_dtypes():
+    dev = dpctl.select_default_device()
+
+    int_dt = default_device_int_type(dev)
+    ind_dt = default_device_index_type(dev)
+    fp_dt = default_device_fp_type(dev)
+    cm_dt = default_device_complex_type(dev)
+
+    info = dpt.__array_namespace_info__()
+    default_dts_nodev = info.default_dtypes()
+    default_dts_dev = info.default_dtypes(dev)
+
+    assert (
+        int_dt == default_dts_nodev["integral"] == default_dts_dev["integral"]
+    )
+    assert (
+        ind_dt == default_dts_nodev["indexing"] == default_dts_dev["indexing"]
+    )
+    assert (
+        fp_dt
+        == default_dts_nodev["real floating"]
+        == default_dts_dev["real floating"]
+    )
+    assert (
+        cm_dt
+        == default_dts_nodev["complex floating"]
+        == default_dts_dev["complex floating"]
+    )
+
+
+def test_array_api_inspection_default_device_dtypes():
+    dev = dpctl.select_default_device()
+    dtypes = _dtypes_no_fp16_fp64.copy()
+    if dev.has_aspect_fp64:
+        dtypes["float64"] = dpt.float64
+
+    assert dtypes == dpt.__array_namespace_info__().dtypes()
+
+
+@pytest.mark.parametrize("fp16", [True, False])
+@pytest.mark.parametrize("fp64", [True, False])
+def test_array_api_inspection_device_dtypes(fp16, fp64):
+    dev = MockDevice(fp16, fp64)
+    dtypes = _dtypes_no_fp16_fp64.copy()
+    if fp64:
+        dtypes["float64"] = dpt.float64
+
+    assert dtypes == dpt.__array_namespace_info__().dtypes(device=dev)
+
+
+def test_array_api_inspection_dtype_kind():
+    info = dpt.__array_namespace_info__()
+
+    f_dtypes = info.dtypes(kind="real floating")
+    assert all([_dt[1].kind == "f" for _dt in f_dtypes.items()])
+
+    i_dtypes = info.dtypes(kind="signed integer")
+    assert all([_dt[1].kind == "i" for _dt in i_dtypes.items()])
+
+    u_dtypes = info.dtypes(kind="unsigned integer")
+    assert all([_dt[1].kind == "u" for _dt in u_dtypes.items()])
+
+    ui_dtypes = info.dtypes(kind="unsigned integer")
+    assert all([_dt[1].kind in "ui" for _dt in ui_dtypes.items()])
+
+    c_dtypes = info.dtypes(kind="complex floating")
+    assert all([_dt[1].kind == "c" for _dt in c_dtypes.items()])
+
+    assert info.dtypes(kind="bool") == {"bool": dpt.bool}
+
+    _signed_ints = {
+        "int8": dpt.int8,
+        "int16": dpt.int16,
+        "int32": dpt.int32,
+        "int64": dpt.int64,
+    }
+    assert (
+        info.dtypes(kind=("signed integer", "signed integer")) == _signed_ints
+    )
+    assert (
+        info.dtypes(
+            kind=("integral", "bool", "real floating", "complex floating")
+        )
+        == info.dtypes()
+    )
