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'