Merge pull request #1793 from IntelPython/document-dldevicetype-constants

Expose `DLDeviceType` enumerator in `dpctl.tensor`

Author: oleksandr-pavlyk

docs/doc_sources/api_reference/dpctl/tensor.constants.rst

@@ -0,0 +1,35 @@
+.. _dpctl_tensor_constants:
+
+Constants
+========================
+
+The following constants are defined in :py:mod:`dpctl.tensor`:
+
+.. currentmodule:: dpctl.tensor
+
+.. autodata:: DLDeviceType
+
+.. data:: e
+
+   ``float``:
+        IEEE 754 floating-point representation of Euler's constant.
+
+.. data:: inf
+
+   ``float``:
+        IEEE 754 floating-point representation of (positive) infinity.
+
+.. data:: nan
+
+   ``float``:
+        IEEE 754 floating-point representation of Not a Number (NaN).
+
+.. data:: newaxis
+
+   ``NoneType``:
+        Alias for ``None`` which is useful for indexing.
+
+.. data:: pi
+
+   ``float``:
+        IEEE 754 floating-point representation of the mathematical constant π.

docs/doc_sources/api_reference/dpctl/tensor.rst

@@ -29,6 +29,7 @@ This module contains:
 * :ref:`sorting functions <dpctl_tensor_sorting_functions>`
 * :ref:`statistical functions <dpctl_tensor_statistical_functions>`
 * :ref:`utility functions <dpctl_tensor_utility_functions>`
+* :ref:`constants <dpctl_tensor_constants>`
 
 
 .. toctree::
@@ -48,3 +49,4 @@ This module contains:
     tensor.sorting_functions
     tensor.statistical_functions
     tensor.utility_functions
+    tensor.constants

dpctl/tensor/__init__.py

@@ -93,7 +93,7 @@
 from dpctl.tensor._reshape import reshape
 from dpctl.tensor._search_functions import where
 from dpctl.tensor._statistical_functions import mean, std, var
-from dpctl.tensor._usmarray import usm_ndarray
+from dpctl.tensor._usmarray import DLDeviceType, usm_ndarray
 from dpctl.tensor._utility_functions import all, any, diff
 
 from ._accumulation import cumulative_logsumexp, cumulative_prod, cumulative_sum
@@ -376,4 +376,5 @@
     "nextafter",
     "diff",
     "count_nonzero",
+    "DLDeviceType",
 ]

dpctl/tensor/_usmarray.pyx

@@ -54,6 +54,40 @@ include "_slicing.pxi"
 
 
 class DLDeviceType(IntEnum):
+    """
+    An ``IntEnum`` for the types of DLDevices supported by the DLPack
+    protocol.
+        ``kDLCPU``:
+            CPU (host) device
+        ``kDLCUDA``:
+            CUDA GPU device
+        ``kDLCUDAHost``:
+            Pinned CUDA CPU memory by cudaMallocHost
+        ``kDLOpenCL``:
+            OpenCL device
+        ``kDLVulkan``:
+            Vulkan buffer
+        ``kDLMetal``:
+            Metal for Apple GPU
+        ``kDLVPI``:
+            Verilog simulator buffer
+        ``kDLROCM``:
+            ROCm GPU device
+        ``kDLROCMHost``:
+            Pinned ROCm CPU memory allocated by hipMallocHost
+        ``kDLExtDev``:
+            Reserved extension device type used to test new devices
+        ``kDLCUDAManaged``:
+            CUDA managed/unified memory allocated by cudaMallocManaged
+        ``kDLOneAPI``:
+            Unified shared memory allocated on a oneAPI non-partitioned device
+        ``kDLWebGPU``:
+            Device support for WebGPU standard
+        ``kDLHexagon``:
+            Qualcomm Hexagon DSP
+        ``kDLMAIA``:
+            Microsoft MAIA device
+    """
     kDLCPU = c_dlpack.device_CPU
     kDLCUDA = c_dlpack.device_CUDA
     kDLCUDAHost = c_dlpack.device_CUDAHost
@@ -1244,6 +1278,8 @@ cdef class usm_ndarray:
         The tuple describes the non-partitioned device where the array has been allocated,
         or the non-partitioned parent device of the allocation device.
 
+        See ``DLDeviceType`` for a list of devices supported by the DLPack protocol.
+
         Raises:
             DLPackCreationError:
                 when the ``device_id`` could not be determined.