Merge branch 'main' into cupy_device

Author: crusaderky

array_api_compat/common/_aliases.py

@@ -5,7 +5,7 @@
 from __future__ import annotations
 
 import inspect
-from typing import NamedTuple, Optional, Sequence, Tuple, Union
+from typing import Any, NamedTuple, Optional, Sequence, Tuple, Union
 
 from ._typing import Array, Device, DType, Namespace
 from ._helpers import (
@@ -609,13 +609,30 @@ def sign(x: Array, /, xp: Namespace, **kwargs) -> Array:
         out[xp.isnan(x)] = xp.nan
     return out[()]
 
+
+def finfo(type_: DType | Array, /, xp: Namespace) -> Any:
+    # It is surprisingly difficult to recognize a dtype apart from an array.
+    # np.int64 is not the same as np.asarray(1).dtype!
+    try:
+        return xp.finfo(type_)
+    except (ValueError, TypeError):
+        return xp.finfo(type_.dtype)
+
+
+def iinfo(type_: DType | Array, /, xp: Namespace) -> Any:
+    try:
+        return xp.iinfo(type_)
+    except (ValueError, TypeError):
+        return xp.iinfo(type_.dtype)
+
+
 __all__ = ['arange', 'empty', 'empty_like', 'eye', 'full', 'full_like',
            'linspace', 'ones', 'ones_like', 'zeros', 'zeros_like',
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'std', 'var', 'cumulative_sum', 'cumulative_prod','clip', 'permute_dims',
            'reshape', 'argsort', 'sort', 'nonzero', 'ceil', 'floor', 'trunc',
            'matmul', 'matrix_transpose', 'tensordot', 'vecdot', 'isdtype',
-           'unstack', 'sign']
+           'unstack', 'sign', 'finfo', 'iinfo']
 
 _all_ignore = ['inspect', 'array_namespace', 'NamedTuple']

array_api_compat/common/_helpers.py

@@ -598,6 +598,44 @@ def your_function(x, y):
 get_namespace = array_namespace
 
 
+def _device_ctx(
+    bare_xp: Namespace, device: Device, like: Array | None = None
+) -> Generator[None]:
+    """Context manager which changes the current device in CuPy.
+
+    Used internally by array creation functions in common._aliases.
+    """
+    if device is None:
+        if like is None:
+            return contextlib.nullcontext()    
+        device = _device(like)
+
+    if bare_xp is sys.modules.get('numpy'):
+        if device != "cpu":
+            raise ValueError(f"Unsupported device for NumPy: {device!r}")
+        return contextlib.nullcontext()
+
+    if bare_xp is sys.modules.get('dask.array'):
+        if device not in ("cpu", _DASK_DEVICE):
+            raise ValueError(f"Unsupported device for Dask: {device!r}")
+        return contextlib.nullcontext()
+
+    if bare_xp is sys.modules.get('cupy'):
+        if not isinstance(device, bare_xp.cuda.Device):
+            raise TypeError(f"device is not a cupy.cuda.Device: {device!r}")
+        return device
+
+    # PyTorch doesn't have a "current device" context manager and you
+    # can't use array creation functions from common._aliases.
+    raise AssertionError("unreachable")  # pragma: nocover
+
+
+def _check_device(bare_xp: Namespace, device: Device) -> None:
+    """Validate dummy device on device-less array backends."""
+    with _device_ctx(bare_xp, device):
+        pass
+
+
 # Placeholder object to represent the dask device
 # when the array backend is not the CPU.
 # (since it is not easy to tell which device a dask array is on)
@@ -607,7 +645,6 @@ def __repr__(self):
 
 _DASK_DEVICE = _dask_device()
 
-
 # device() is not on numpy.ndarray or dask.array and to_device() is not on numpy.ndarray
 # or cupy.ndarray. They are not included in array objects of this library
 # because this library just reuses the respective ndarray classes without
@@ -799,43 +836,6 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
     return x.to_device(device, stream=stream)
 
 
-def _device_ctx(
-    bare_xp: Namespace, device: Device, like: Array | None = None
-) -> Generator[None]:
-    """Context manager which changes the current device in CuPy.
-
-    Used internally by array creation functions in common._aliases.
-    """
-    if device is None:
-        if like is None:
-            return contextlib.nullcontext()    
-        device = _device(like)
-
-    if bare_xp is sys.modules.get('numpy'):
-        if device != "cpu":
-            raise ValueError(f"Unsupported device for NumPy: {device!r}")
-        return contextlib.nullcontext()
-
-    if bare_xp is sys.modules.get('dask.array'):
-        if device not in ("cpu", _DASK_DEVICE):
-            raise ValueError(f"Unsupported device for Dask: {device!r}")
-        return contextlib.nullcontext()
-
-    if bare_xp is sys.modules.get('cupy'):
-        if not isinstance(device, bare_xp.cuda.Device):
-            raise TypeError(f"device is not a cupy.cuda.Device: {device!r}")
-        return device
-
-    # PyTorch doesn't have a "current device" context manager and you
-    # can't use array creation functions from common._aliases.
-    raise AssertionError("unreachable")  # pragma: nocover
-
-
-def _check_device(bare_xp: Namespace, device: Device) -> None:
-    with _device_ctx(bare_xp, device):
-        pass
-
-
 def size(x: Array) -> int | None:
     """
     Return the total number of elements of x.

array_api_compat/common/_linalg.py

@@ -174,3 +174,5 @@ def trace(
            'svd', 'cholesky', 'matrix_rank', 'pinv', 'matrix_norm',
            'matrix_transpose', 'svdvals', 'vecdot', 'vector_norm', 'diagonal',
            'trace']
+
+_all_ignore = ['math', 'normalize_axis_tuple', 'get_xp', 'np', 'isdtype']

array_api_compat/cupy/__init__.py

@@ -8,9 +8,6 @@
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
-
 __import__(__package__ + '.fft')
 
-from ..common._helpers import * # noqa: F401,F403
-
 __array_api_version__ = '2024.12'

array_api_compat/cupy/_aliases.py

@@ -61,6 +61,8 @@
 matrix_transpose = get_xp(cp)(_aliases.matrix_transpose)
 tensordot = get_xp(cp)(_aliases.tensordot)
 sign = get_xp(cp)(_aliases.sign)
+finfo = get_xp(cp)(_aliases.finfo)
+iinfo = get_xp(cp)(_aliases.iinfo)
 
 
 # asarray also adds the copy keyword, which is not present in numpy 1.0.
@@ -87,7 +89,7 @@ def asarray(
     if copy is False:
         raise NotImplementedError("asarray(copy=False) is not yet supported in cupy")
 
-    like = obj if _helpers.is_cupy_array(obj) else None
+    like = obj if isinstance(obj, cp.ndarray) else None
     with _helpers._device_ctx(cp, device, like=like):
         if copy is None:
             return cp.asarray(obj, dtype=dtype, **kwargs)

array_api_compat/dask/array/__init__.py

@@ -5,5 +5,6 @@
 
 __array_api_version__ = '2024.12'
 
+# See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')

array_api_compat/dask/array/_aliases.py

@@ -5,8 +5,6 @@
 import numpy as np
 from numpy import (
     # dtypes
-    iinfo,
-    finfo,
     bool_ as bool,
     float32,
     float64,
@@ -133,6 +131,8 @@ def arange(
 matmul = get_xp(np)(_aliases.matmul)
 tensordot = get_xp(np)(_aliases.tensordot)
 sign = get_xp(np)(_aliases.sign)
+finfo = get_xp(np)(_aliases.finfo)
+iinfo = get_xp(np)(_aliases.iinfo)
 
 
 # asarray also adds the copy keyword, which is not present in numpy 1.0.
@@ -346,10 +346,9 @@ def count_nonzero(
                     '__array_namespace_info__', 'asarray', 'astype', 'acos',
                     'acosh', 'asin', 'asinh', 'atan', 'atan2',
                     'atanh', 'bitwise_left_shift', 'bitwise_invert',
-                    'bitwise_right_shift', 'concat', 'pow', 'iinfo', 'finfo', 'can_cast',
+                    'bitwise_right_shift', 'concat', 'pow', 'can_cast',
                     'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64',
-                    'uint8', 'uint16', 'uint32', 'uint64',
-                    'complex64', 'complex128', 'iinfo', 'finfo',
+                    'uint8', 'uint16', 'uint32', 'uint64', 'complex64', 'complex128',
                     'can_cast', 'count_nonzero', 'result_type']
 
 _all_ignore = ["array_namespace", "get_xp", "da", "np"]

array_api_compat/numpy/__init__.py

@@ -14,17 +14,8 @@
 # It doesn't overwrite np.linalg from above. The import is generated
 # dynamically so that the library can be vendored.
 __import__(__package__ + '.linalg')
-
 __import__(__package__ + '.fft')
 
 from .linalg import matrix_transpose, vecdot # noqa: F401
 
-from ..common._helpers import * # noqa: F403
-
-try:
-    # Used in asarray(). Not present in older versions.
-    from numpy import _CopyMode # noqa: F401
-except ImportError:
-    pass
-
 __array_api_version__ = '2024.12'

array_api_compat/numpy/_aliases.py

@@ -61,6 +61,8 @@
 matrix_transpose = get_xp(np)(_aliases.matrix_transpose)
 tensordot = get_xp(np)(_aliases.tensordot)
 sign = get_xp(np)(_aliases.sign)
+finfo = get_xp(np)(_aliases.finfo)
+iinfo = get_xp(np)(_aliases.iinfo)
 
 
 def _supports_buffer_protocol(obj):
@@ -86,7 +88,7 @@ def asarray(
     *,
     dtype: Optional[DType] = None,
     device: Optional[Device] = None,
-    copy: "Optional[Union[bool, np._CopyMode]]" = None,
+    copy: Optional[Union[bool, np._CopyMode]] = None,
     **kwargs,
 ) -> Array:
     """

array_api_compat/torch/__init__.py

@@ -9,16 +9,14 @@
         or 'cpu' in n
         or 'backward' in n):
         continue
-    exec(n + ' = torch.' + n)
+    exec(f"{n} = torch.{n}")
+del n
 
 # These imports may overwrite names from the import * above.
 from ._aliases import * # noqa: F403
 
 # See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
-
 __import__(__package__ + '.fft')
 
-from ..common._helpers import * # noqa: F403
-
 __array_api_version__ = '2024.12'

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'.
 
@@ -227,6 +250,9 @@ def min(x: Array, /, *, axis: Optional[Union[int, Tuple[int, ...]]] = None, keep
 unstack = get_xp(torch)(_aliases.unstack)
 cumulative_sum = get_xp(torch)(_aliases.cumulative_sum)
 cumulative_prod = get_xp(torch)(_aliases.cumulative_prod)
+finfo = get_xp(torch)(_aliases.finfo)
+iinfo = get_xp(torch)(_aliases.iinfo)
+
 
 # torch.sort also returns a tuple
 # https://github.com/pytorch/pytorch/issues/70921
@@ -282,7 +308,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 +343,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 +372,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 +396,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 +838,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',
@@ -832,6 +854,6 @@ def sign(x: Array, /) -> Array:
            'UniqueAllResult', 'UniqueCountsResult', 'UniqueInverseResult',
            'unique_all', 'unique_counts', 'unique_inverse', 'unique_values',
            'matmul', 'matrix_transpose', 'vecdot', 'tensordot', 'isdtype',
-           'take', 'take_along_axis', 'sign']
+           'take', 'take_along_axis', 'sign', 'finfo', 'iinfo']
 
 _all_ignore = ['torch', 'get_xp']

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

cupy-xfails.txt

@@ -14,9 +14,10 @@ array_api_tests/test_array_object.py::test_getitem
 # copy=False is not yet implemented
 array_api_tests/test_creation_functions.py::test_asarray_arrays
 
-# finfo test is testing that the result is a float instead of float32 (see
-# also https://github.com/data-apis/array-api/issues/405)
+# attributes are np.float32 instead of float
+# (see also https://github.com/data-apis/array-api/issues/405)
 array_api_tests/test_data_type_functions.py::test_finfo[float32]
+array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 
 # Some array attributes are missing, and we do not wrap the array object
 array_api_tests/test_has_names.py::test_has_names[array_method-__array_namespace__]

dask-xfails.txt

@@ -12,8 +12,10 @@ array_api_tests/test_array_object.py::test_getitem_masking
 # zero division error, and typeerror: tuple indices must be integers or slices not tuple
 array_api_tests/test_creation_functions.py::test_eye
 
-# finfo(float32).eps returns float32 but should return float
+# attributes are np.float32 instead of float
+# (see also https://github.com/data-apis/array-api/issues/405)
 array_api_tests/test_data_type_functions.py::test_finfo[float32]
+array_api_tests/test_data_type_functions.py::test_finfo[complex64]
 
 # out[-1]=dask.array<getitem ...> but should be some floating number
 # (I think the test is not forcing the op to be computed?)