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

Author: ndgrigorian

dpctl/tensor/_array_api.py

@@ -23,14 +23,9 @@
     default_device_int_type,
 )
 
-__array_api_version__ = "2022.12"
-
 
 def _isdtype_impl(dtype, kind):
-    if isinstance(kind, dpt.dtype):
-        return dtype == kind
-
-    elif isinstance(kind, str):
+    if isinstance(kind, str):
         if kind == "bool":
             return dtype.kind == "b"
         elif kind == "signed integer":
@@ -54,15 +49,21 @@ def _isdtype_impl(dtype, kind):
         raise TypeError(f"Unsupported data type kind: {kind}")
 
 
-class __array_namespace_info__:
+__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,
-            "float16": dpt.float16,
             "float32": dpt.float32,
             "float64": dpt.float64,
             "complex64": dpt.complex64,
@@ -78,41 +79,129 @@ def __init__(self):
         }
 
     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": default_device_fp_type(device),
-            "complex floating": default_device_complex_type,
-            "integral": default_device_int_type(device),
-            "indexing": default_device_index_type(device),
+            "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()
-        ignored_types = []
-        if not device.has_aspect_fp16:
-            ignored_types.append("float16")
-        if not device.has_aspect_fp64:
-            ignored_types.append("float64")
+        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 not in ignored_types
+                if (key != "float64" or _fp64)
             }
         else:
             return {
                 key: val
                 for key, val in self._all_dtypes.items()
-                if key not in ignored_types and _isdtype_impl(val, kind)
+                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()