MAINT: clarify default_device output

Author: crusaderky

array_api_compat/common/_aliases.py

@@ -18,7 +18,7 @@
 
 # These functions are modified from the NumPy versions.
 
-# Creation functions add the device keyword (which does nothing for NumPy)
+# Creation functions add the device keyword (which does nothing for NumPy and Dask)
 
 def arange(
     start: Union[int, float],

array_api_compat/cupy/_info.py

@@ -26,6 +26,7 @@
     complex128,
 )
 
+
 class __array_namespace_info__:
     """
     Get the array API inspection namespace for CuPy.
@@ -117,7 +118,7 @@ def default_device(self):
 
         Returns
         -------
-        device : str
+        device : Device
             The default device used for new CuPy arrays.
 
         Examples
@@ -126,6 +127,15 @@ def default_device(self):
         >>> info.default_device()
         Device(0)
 
+        Notes
+        -----
+        This method returns the static default device when CuPy is initialized.
+        However, the *current* device used by creation functions (``empty`` etc.)
+        can be changed globally or with a context manager.
+
+        See Also
+        --------
+        https://github.com/data-apis/array-api/issues/835
         """
         return cuda.Device(0)
 
@@ -312,7 +322,7 @@ def devices(self):
 
         Returns
         -------
-        devices : list of str
+        devices : list[Device]
             The devices supported by CuPy.
 
         See Also

array_api_compat/dask/array/_info.py

@@ -130,7 +130,7 @@ def default_device(self):
 
         Returns
         -------
-        device : str
+        device : Device
             The default device used for new Dask arrays.
 
         Examples
@@ -335,7 +335,7 @@ def devices(self):
 
         Returns
         -------
-        devices : list of str
+        devices : list[Device]
             The devices supported by Dask.
 
         See Also

array_api_compat/numpy/_info.py

@@ -119,7 +119,7 @@ def default_device(self):
 
         Returns
         -------
-        device : str
+        device : Device
             The default device used for new NumPy arrays.
 
         Examples
@@ -326,7 +326,7 @@ def devices(self):
 
         Returns
         -------
-        devices : list of str
+        devices : list[Device]
             The devices supported by NumPy.
 
         See Also

array_api_compat/torch/_info.py

@@ -102,15 +102,24 @@ def default_device(self):
 
         Returns
         -------
-        device : str
+        device : Device
             The default device used for new PyTorch arrays.
 
         Examples
         --------
         >>> info = np.__array_namespace_info__()
         >>> info.default_device()
-        'cpu'
+        device(type='cpu')
 
+        Notes
+        -----
+        This method returns the static default device when PyTorch is initialized.
+        However, the *current* device used by creation functions (``empty`` etc.)
+        can be changed at runtime.
+
+        See Also
+        --------
+        https://github.com/data-apis/array-api/issues/835
         """
         return torch.device("cpu")
 
@@ -120,9 +129,9 @@ def default_dtypes(self, *, device=None):
 
         Parameters
         ----------
-        device : str, optional
-            The device to get the default data types for. For PyTorch, only
-            ``'cpu'`` is allowed.
+        device : Device, optional
+            The device to get the default data types for.
+            Unused for PyTorch, as all devices use the same default dtypes.
 
         Returns
         -------
@@ -250,8 +259,9 @@ def dtypes(self, *, device=None, kind=None):
 
         Parameters
         ----------
-        device : str, optional
+        device : Device, optional
             The device to get the data types for.
+            Unused for PyTorch, as all devices use the same dtypes.
         kind : str or tuple of str, optional
             The kind of data types to return. If ``None``, all data types are
             returned. If a string, only data types of that kind are returned.
@@ -310,7 +320,7 @@ def devices(self):
 
         Returns
         -------
-        devices : list of str
+        devices : list[Device]
             The devices supported by PyTorch.
 
         See Also
@@ -333,6 +343,7 @@ def devices(self):
         # device:
         try:
             torch.device('notadevice')
+            raise AssertionError("unreachable")  # pragma: nocover
         except RuntimeError as e:
             # The error message is something like:
             # "Expected one of cpu, cuda, ipu, xpu, mkldnn, opengl, opencl, ideep, hip, ve, fpga, ort, xla, lazy, vulkan, mps, meta, hpu, mtia, privateuseone device type at start of device string: notadevice"