diff --git a/.github/workflows/conda-package.yml b/.github/workflows/conda-package.yml index b09be78b08..b6a4649a30 100644 --- 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 index 5eee3e9ab9..8638fc6d29 100644 --- 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 new file mode 100644 index 0000000000..613d6dcd66 --- /dev/null +++ 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 new file mode 100644 index 0000000000..5ae0d35f8e --- /dev/null +++ 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() + )