ENH: torch.asarray device propagation

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/_aliases.py

@@ -2,12 +2,13 @@
 
 from functools import reduce as _reduce, wraps as _wraps
 from builtins import all as _builtin_all, any as _builtin_any
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import Any, List, Optional, Sequence, Tuple, Union
 
 import torch
 
 from .._internal import get_xp
 from ..common import _aliases
+from ..common._typing import NestedSequence, SupportsBufferProtocol
 from ._info import __array_namespace_info__
 from ._typing import Array, Device, DType
 
@@ -207,6 +208,28 @@ def can_cast(from_: Union[DType, Array], to: DType, /) -> bool:
 remainder = _two_arg(torch.remainder)
 subtract = _two_arg(torch.subtract)
 
+
+def asarray(
+    obj: (
+    Array 
+        | bool | int | float | complex 
+        | NestedSequence[bool | int | float | complex] 
+        | SupportsBufferProtocol
+    ),
+    /,
+    *,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    copy: bool | None = None,
+    **kwargs: Any,
+) -> Array:
+    # torch.asarray does not respect input->output device propagation
+    # https://github.com/pytorch/pytorch/issues/150199
+    if device is None and isinstance(obj, torch.Tensor):
+        device = obj.device
+    return torch.asarray(obj, dtype=dtype, device=device, copy=copy, **kwargs)
+
+
 # These wrappers are mostly based on the fact that pytorch uses 'dim' instead
 # of 'axis'.
 
@@ -282,7 +305,6 @@ def prod(x: Array,
          dtype: Optional[DType] = None,
          keepdims: bool = False,
          **kwargs) -> Array:
-    x = torch.asarray(x)
     ndim = x.ndim
 
     # https://github.com/pytorch/pytorch/issues/29137. Separate from the logic
@@ -318,7 +340,6 @@ def sum(x: Array,
          dtype: Optional[DType] = None,
          keepdims: bool = False,
          **kwargs) -> Array:
-    x = torch.asarray(x)
     ndim = x.ndim
 
     # https://github.com/pytorch/pytorch/issues/29137.
@@ -348,7 +369,6 @@ def any(x: Array,
         axis: Optional[Union[int, Tuple[int, ...]]] = None,
         keepdims: bool = False,
         **kwargs) -> Array:
-    x = torch.asarray(x)
     ndim = x.ndim
     if axis == ():
         return x.to(torch.bool)
@@ -373,7 +393,6 @@ def all(x: Array,
         axis: Optional[Union[int, Tuple[int, ...]]] = None,
         keepdims: bool = False,
         **kwargs) -> Array:
-    x = torch.asarray(x)
     ndim = x.ndim
     if axis == ():
         return x.to(torch.bool)
@@ -816,7 +835,7 @@ def sign(x: Array, /) -> Array:
         return out
 
 
-__all__ = ['__array_namespace_info__', 'result_type', 'can_cast',
+__all__ = ['__array_namespace_info__', 'asarray', 'result_type', 'can_cast',
            'permute_dims', 'bitwise_invert', 'newaxis', 'conj', 'add',
            'atan2', 'bitwise_and', 'bitwise_left_shift', 'bitwise_or',
            'bitwise_right_shift', 'bitwise_xor', 'copysign', 'count_nonzero',

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"

array_api_compat/torch/_typing.py

@@ -1,4 +1,3 @@
-__all__ = ["Array", "DType", "Device"]
+__all__ = ["Array", "Device", "DType"]
 
-from torch import dtype as DType, Tensor as Array
-from ..common._typing import Device
+from torch import device as Device, dtype as DType, Tensor as Array