array-api-strict 2.4.1__py3-none-any.whl → 2.6__py3-none-any.whl

array_api_strict/__init__.py

@@ -73,6 +73,7 @@ __all__ += [
 from ._data_type_functions import (
     astype,
     broadcast_arrays,
+    broadcast_shapes,
     broadcast_to,
     can_cast,
     finfo,
@@ -84,6 +85,7 @@ from ._data_type_functions import (
 __all__ += [
     "astype",
     "broadcast_arrays",
+    "broadcast_shapes",
     "broadcast_to",
     "can_cast",
     "finfo",
@@ -299,9 +301,9 @@ from ._searching_functions import argmax, argmin, nonzero, count_nonzero, search
 
 __all__ += ["argmax", "argmin", "nonzero", "count_nonzero", "searchsorted", "where"]
 
-from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values
+from ._set_functions import unique_all, unique_counts, unique_inverse, unique_values, isin
 
-__all__ += ["unique_all", "unique_counts", "unique_inverse", "unique_values"]
+__all__ += ["unique_all", "unique_counts", "unique_inverse", "unique_values", "isin"]
 
 from ._sorting_functions import argsort, sort
 

array_api_strict/_array_object.py

@@ -20,7 +20,7 @@ import sys
 from collections.abc import Iterator
 from enum import IntEnum
 from types import EllipsisType, ModuleType
-from typing import Any, Final, Literal, SupportsIndex
+from typing import Any, Literal, SupportsIndex, Callable
 
 import numpy as np
 import numpy.typing as npt
@@ -40,35 +40,14 @@ from ._dtypes import (
     _real_to_complex_map,
     _result_type,
 )
+from ._devices import (
+    CPU_DEVICE, Device, device_supports_dtype, _normalize_dl_device, _DLPACK_DEVICE_FOR,
+    DLDeviceType
+)
 from ._flags import get_array_api_strict_flags, set_array_api_strict_flags
 from ._typing import PyCapsule
 
 
-class Device:
-    _device: Final[str]
-    __slots__ = ("_device", "__weakref__")
-
-    def __init__(self, device: str = "CPU_DEVICE"):
-        if device not in ("CPU_DEVICE", "device1", "device2"):
-            raise ValueError(f"The device '{device}' is not a valid choice.")
-        self._device = device
-
-    def __repr__(self) -> str:
-        return f"array_api_strict.Device('{self._device}')"
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Device):
-            return False
-        return self._device == other._device
-
-    def __hash__(self) -> int:
-        return hash(("Device", self._device))
-
-
-CPU_DEVICE = Device()
-ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"))
-
-
 class Array:
     """
     n-d array object for the array API namespace.
@@ -113,10 +92,15 @@ class Array:
             raise TypeError(
                 f"The array_api_strict namespace does not support the dtype '{x.dtype}'"
             )
-        obj._array = x
-        obj._dtype = _dtype
+
         if device is None:
             device = CPU_DEVICE
+        if not device_supports_dtype(device, _dtype):
+            raise ValueError(f"Device {device!r} does not support dtype={_dtype!r}.")
+
+        obj._array = x
+        obj._dtype = _dtype
+
         obj._device = device
         return obj
 
@@ -125,6 +109,9 @@ class Array:
         raise TypeError(
             "The array_api_strict Array object should not be instantiated directly. Use an array creation function, such as asarray(), instead."
         )
+    
+    def __reduce__(self) -> tuple[Callable, tuple[npt.NDArray[Any], Device]]:
+        return (self._new, (self._array, self._device))
 
     # These functions are not required by the spec, but are implemented for
     # the sake of usability.
@@ -397,7 +384,7 @@ class Array:
                 isinstance(i, SupportsIndex)  # i.e. ints
                 or isinstance(i, slice)
                 or i == Ellipsis
-                or i is None
+                or (op == "getitem" and i is None) # `None` disallowed in setitem
                 or isinstance(i, Array)
                 or isinstance(i, np.ndarray)
             ):
@@ -627,22 +614,40 @@ class Array:
                 raise NotImplementedError("The copy argument to __dlpack__ is not yet implemented")
 
             return self._array.__dlpack__(stream=stream)
-        else:
-            kwargs = {'stream': stream}
-            if max_version is not _undef:
-                kwargs['max_version'] = max_version
-            if dl_device is not _undef:
-                kwargs['dl_device'] = dl_device
-            if copy is not _undef:
-                kwargs['copy'] = copy
-            return self._array.__dlpack__(**kwargs)
+
+        kwargs: dict[str, Any] = {'stream': stream}
+        self_dl_device = _normalize_dl_device(*_DLPACK_DEVICE_FOR[self._device])
+        cpu_dl_device = _normalize_dl_device(DLDeviceType.kDLCPU, 0)
+        numpy_dl_device: tuple[IntEnum, int] | None = None
+
+        if dl_device not in [_undef, None]:
+            requested = _normalize_dl_device(dl_device[0], dl_device[1])
+            if requested == self_dl_device:
+                pass
+            elif requested == cpu_dl_device and self_dl_device != cpu_dl_device:
+                if copy is False:
+                    raise BufferError(
+                        "Cannot export array to CPU without copying when copy=False"
+                    )
+                if copy is _undef:
+                    copy = True
+                numpy_dl_device = (DLDeviceType.kDLCPU, 0)
+            else:
+                raise BufferError("unsupported device requested")
+
+        if max_version is not _undef:
+            kwargs['max_version'] = max_version
+        if numpy_dl_device is not None:
+            kwargs['dl_device'] = numpy_dl_device
+        if copy is not _undef:
+            kwargs['copy'] = copy
+        return self._array.__dlpack__(**kwargs)
 
     def __dlpack_device__(self) -> tuple[IntEnum, int]:
         """
         Performs the operation __dlpack_device__.
         """
-        # Note: device support is required for this
-        return self._array.__dlpack_device__()
+        return _DLPACK_DEVICE_FOR[self._device]
 
     def __eq__(self, other: Array | complex, /) -> Array:  # type: ignore[override]
         """
@@ -957,6 +962,10 @@ class Array:
         other = value
         if isinstance(value, (bool, int, float, complex)):
             other = self._promote_scalar(value)
+        else:
+            if value.device != self.device:
+                raise ValueError(f"mismatched devices: {self.device = } != {value.device =}.")
+
         dt = _result_type(self.dtype, other.dtype)
         if dt != self.dtype:
             raise TypeError(f"mismatched dtypes: {self.dtype = } and {other.dtype = }")

array_api_strict/_creation_functions.py

@@ -5,7 +5,11 @@ from typing import TYPE_CHECKING, Literal
 
 import numpy as np
 
-from ._dtypes import DType, _all_dtypes, _np_dtype
+from ._dtypes import DType, _all_dtypes, _np_dtype, bool as xp_bool
+from ._devices import (
+    Device, device_supports_dtype, get_default_dtypes,
+    check_device as _check_device
+)
 from ._flags import get_array_api_strict_flags
 from ._typing import NestedSequence, SupportsBufferProtocol, SupportsDLPack
 
@@ -14,7 +18,7 @@ if TYPE_CHECKING:
     from typing_extensions import TypeIs
 
     # Circular import
-    from ._array_object import Array, Device
+    from ._array_object import Array
 
 
 class Undef(Enum):
@@ -24,10 +28,15 @@ class Undef(Enum):
 _undef = Undef.UNDEF
 
 
-def _check_valid_dtype(dtype: DType | None) -> None:
+def _check_valid_dtype(dtype: DType | None, device: Device | None = None) -> None:
     # Note: Only spelling dtypes as the dtype objects is supported.
-    if dtype not in (None,) + _all_dtypes:
-        raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
+    if dtype is not None:
+        if dtype not in _all_dtypes:
+            raise ValueError(f"dtype must be one of the supported dtypes, got {dtype!r}")
+
+        if device is not None:
+            if not device_supports_dtype(device, dtype):
+                raise ValueError(f"Device {device!r} does not support dtype={dtype!r}.")
 
 
 def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
@@ -38,18 +47,6 @@ def _supports_buffer_protocol(obj: object) -> TypeIs[SupportsBufferProtocol]:
     return True
 
 
-def _check_device(device: Device | None) -> None:
-    # _array_object imports in this file are inside the functions to avoid
-    # circular imports
-    from ._array_object import ALL_DEVICES, Device
-
-    if device is not None and not isinstance(device, Device):
-        raise ValueError(f"Unsupported device {device!r}")
-
-    if device is not None and device not in ALL_DEVICES:
-        raise ValueError(f"Unsupported device {device!r}")
-
-
 def asarray(
     obj: Array | complex | NestedSequence[complex] | SupportsBufferProtocol,
     /,
@@ -65,11 +62,12 @@ def asarray(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
+    _check_device(device)
+    _check_valid_dtype(dtype, device)
     _np_dtype = None
     if dtype is not None:
         _np_dtype = dtype._np_dtype
-    _check_device(device)
+
     if isinstance(obj, Array) and device is None:
         device = obj.device
 
@@ -108,6 +106,27 @@ def asarray(
         raise OverflowError("Integer out of bounds for array dtypes")
 
     res = np.array(obj, dtype=_np_dtype, copy=copy)
+
+    # numpy default dtype may differ; if so, adjust the dtype
+    if dtype is None and device is not None:
+        res_dtype = DType(res.dtype)
+        # The dtype selected by Numpy might not be the default dtype
+        # on this device. We thus find the default dtype for the dtype "kind", and
+        # cast to the device-appropriate default.
+        from ._data_type_functions import isdtype
+        if isdtype(res_dtype, "bool"):
+            target_dtype = DType("bool")
+        elif isdtype(res_dtype, "integral"):
+            target_dtype = get_default_dtypes(device)["integral"]
+        elif isdtype(res_dtype, "real floating"):
+            target_dtype = get_default_dtypes(device)["real floating"]
+        elif isdtype(res_dtype, "complex floating"):
+            target_dtype = get_default_dtypes(device)["complex floating"]
+        else:
+            raise ValueError(f"{res_dtype = } not understood.")
+
+        res = res.astype(target_dtype._np_dtype)
+
     return Array._new(res, device=device)
 
 
@@ -127,8 +146,13 @@ def arange(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        if any(isinstance(x, float) for x in (start, stop, step)):
+            dtype = get_default_dtypes(device)["real floating"]
+        else:
+            dtype = get_default_dtypes(device)["integral"]
 
     return Array._new(
         np.arange(start, stop, step, dtype=_np_dtype(dtype)),
@@ -149,8 +173,10 @@ def empty(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        dtype = get_default_dtypes(device)["real floating"]
 
     return Array._new(np.empty(shape, dtype=_np_dtype(dtype)), device=device)
 
@@ -165,10 +191,12 @@ def empty_like(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
     if device is None:
         device = x.device
+    if dtype is None:
+        dtype = x.dtype
+    _check_valid_dtype(dtype, device)
 
     return Array._new(np.empty_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
@@ -189,8 +217,10 @@ def eye(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        dtype = get_default_dtypes(device)["real floating"]
 
     return Array._new(
         np.eye(n_rows, M=n_cols, k=k, dtype=_np_dtype(dtype)), device=device
@@ -212,15 +242,20 @@ def from_dlpack(
         if copy is not _undef:
             raise ValueError("The copy argument to from_dlpack requires at least version 2023.12 of the array API")
 
-    # Going to wait for upstream numpy support
     if device is not _undef:
         _check_device(device)
     else:
         device = None
-    if copy not in [_undef, None]:
-        raise NotImplementedError("The copy argument to from_dlpack is not yet implemented")
+        if hasattr(x, "__dlpack_device__"):
+            from ._devices import _device_from_dlpack_device
 
-    return Array._new(np.from_dlpack(x), device=device)
+            dl_type, dl_id = x.__dlpack_device__()
+            device = _device_from_dlpack_device(dl_type, dl_id)
+    if copy in [_undef, None]:
+        # numpy 1.26 does not have the copy= arg
+        return Array._new(np.from_dlpack(x), device=device)
+
+    return Array._new(np.from_dlpack(x, copy=copy), device=device)
 
 
 def full(
@@ -237,11 +272,22 @@ def full(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+
+    if not isinstance(fill_value, bool | int | float | complex):
+        msg = f"Expected Python scalar fill_value, got type {type(fill_value)}"
+        raise TypeError(msg)
+
+    if dtype is None:
+        if type(fill_value) == bool:
+            dtype = xp_bool
+        else:
+            kind = {
+                int: "integral", float: "real floating", complex: "complex floating"
+            }[type(fill_value)]
+            dtype = get_default_dtypes(device)[kind]
 
-    if isinstance(fill_value, Array) and fill_value.ndim == 0:
-        fill_value = fill_value._array
     res = np.full(shape, fill_value, dtype=_np_dtype(dtype))
     if DType(res.dtype) not in _all_dtypes:
         # This will happen if the fill value is not something that NumPy
@@ -265,10 +311,16 @@ def full_like(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
     if device is None:
         device = x.device
+    if dtype is None:
+        dtype = x.dtype
+    _check_valid_dtype(dtype, device)
+
+    if not isinstance(fill_value, bool | int | float | complex):
+        msg = f"Expected Python scalar fill_value, got type {type(fill_value)}"
+        raise TypeError(msg)
 
     res = np.full_like(x._array, fill_value, dtype=_np_dtype(dtype))
     if DType(res.dtype) not in _all_dtypes:
@@ -295,8 +347,13 @@ def linspace(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        if isinstance(start, complex) or isinstance(stop, complex):
+            dtype = get_default_dtypes(device)["complex floating"]
+        else:
+            dtype = get_default_dtypes(device)["real floating"]
 
     return Array._new(
         np.linspace(start, stop, num, dtype=_np_dtype(dtype), endpoint=endpoint),
@@ -304,7 +361,7 @@ def linspace(
     )
 
 
-def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> list[Array]:
+def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> tuple[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.meshgrid <numpy.meshgrid>`.
 
@@ -327,10 +384,12 @@ def meshgrid(*arrays: Array, indexing: Literal["xy", "ij"] = "xy") -> list[Array
     else:
         device = None
 
-    return [
+    typ = list if get_array_api_strict_flags()['api_version'] < '2025.12' else tuple
+
+    return typ(
         Array._new(array, device=device)
         for array in np.meshgrid(*[a._array for a in arrays], indexing=indexing)
-    ]
+    )
 
 
 def ones(
@@ -346,8 +405,10 @@ def ones(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        dtype = get_default_dtypes(device)["real floating"]
 
     return Array._new(np.ones(shape, dtype=_np_dtype(dtype)), device=device)
 
@@ -362,10 +423,12 @@ def ones_like(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
     if device is None:
         device = x.device
+    if dtype is None:
+        dtype = x.dtype
+    _check_valid_dtype(dtype, device)
 
     return Array._new(np.ones_like(x._array, dtype=_np_dtype(dtype)), device=device)
 
@@ -411,8 +474,10 @@ def zeros(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
+    _check_valid_dtype(dtype, device)
+    if dtype is None:
+        dtype = get_default_dtypes(device)["real floating"]
 
     return Array._new(np.zeros(shape, dtype=_np_dtype(dtype)), device=device)
 
@@ -427,9 +492,11 @@ def zeros_like(
     """
     from ._array_object import Array
 
-    _check_valid_dtype(dtype)
     _check_device(device)
     if device is None:
         device = x.device
+    if dtype is None:
+        dtype = x.dtype
+    _check_valid_dtype(dtype, device)
 
     return Array._new(np.zeros_like(x._array, dtype=_np_dtype(dtype)), device=device)

array_api_strict/_data_type_functions.py

@@ -16,7 +16,7 @@ from ._dtypes import (
     _signed_integer_dtypes,
     _unsigned_integer_dtypes,
 )
-from ._flags import get_array_api_strict_flags
+from ._flags import get_array_api_strict_flags, requires_api_version
 
 
 # Note: astype is a function, not an array method as in NumPy.
@@ -49,7 +49,7 @@ def astype(
     return Array._new(x._array.astype(dtype=dtype._np_dtype, copy=copy), device=device)
 
 
-def broadcast_arrays(*arrays: Array) -> list[Array]:
+def broadcast_arrays(*arrays: Array) -> tuple[Array, ...]:
     """
     Array API compatible wrapper for :py:func:`np.broadcast_arrays <numpy.broadcast_arrays>`.
 
@@ -57,9 +57,21 @@ def broadcast_arrays(*arrays: Array) -> list[Array]:
     """
     from ._array_object import Array
 
-    return [
+    typ = list if get_array_api_strict_flags()['api_version'] < '2025.12' else tuple
+
+    return typ(
         Array._new(array, device=arrays[0].device) for array in np.broadcast_arrays(*[a._array for a in arrays])
-    ]
+    )
+
+
+@requires_api_version("2025.12")
+def broadcast_shapes(*shapes: tuple[int, ...]) -> tuple[int, ...]:
+    """
+    Array API compatible wrapper for :py:func:`np.broadcast_shapes <numpy.broadcast_shapes>`.
+
+    See its docstring for more information.
+    """
+    return np.broadcast_shapes(*shapes)
 
 
 def broadcast_to(x: Array, /, shape: tuple[int, ...]) -> Array:

array_api_strict/_devices.py

@@ -0,0 +1,136 @@
+from typing import Final
+from enum import IntEnum
+
+from ._dtypes import (
+    DType, float32, float64, complex64, complex128, int64,
+    _all_dtypes, _boolean_dtypes, _signed_integer_dtypes,
+    _unsigned_integer_dtypes, _integer_dtypes, _real_floating_dtypes,
+    _complex_floating_dtypes, _numeric_dtypes
+)
+
+_ALL_DEVICE_NAMES = ("CPU_DEVICE", "device1", "device2", "no_float64")
+
+class Device:
+    _device: Final[str]
+    __slots__ = ("_device", "__weakref__")
+
+    def __init__(self, device: str = "CPU_DEVICE"):
+        if device not in _ALL_DEVICE_NAMES:
+            raise ValueError(f"The device '{device}' is not a valid choice.")
+        self._device = device
+
+    def __repr__(self) -> str:
+        return f"array_api_strict.Device('{self._device}')"
+
+    def __eq__(self, other: object) -> bool:
+        if not isinstance(other, Device):
+            return False
+        return self._device == other._device
+
+    def __hash__(self) -> int:
+        return hash(("Device", self._device))
+
+
+CPU_DEVICE = Device()
+NO_FLOAT64_DEVICE = Device("no_float64")
+
+ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE)
+
+class DLDeviceType(IntEnum):
+    kDLCPU = 1
+    kDLCUDA = 2
+    kDLMETAL = 8
+
+
+_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = {
+    CPU_DEVICE: (DLDeviceType.kDLCPU, 0),
+    Device("device1"): (DLDeviceType.kDLCUDA, 0),
+    Device("device2"): (DLDeviceType.kDLCUDA, 1),
+    NO_FLOAT64_DEVICE: (DLDeviceType.kDLMETAL, 0),
+}
+
+_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = {
+    (int(device_type), device_id): logical_device
+    for logical_device, (device_type, device_id) in _DLPACK_DEVICE_FOR.items()
+}
+
+
+def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]:
+    return (int(device_type), device_id)
+
+
+def _device_from_dlpack_device(
+    device_type: IntEnum | int, device_id: int
+) -> Device:
+    # NB: if the (device_type, device_id) pair not known, raise
+    return _DLPACK_DEVICE_TO_LOGICAL[_normalize_dl_device(device_type, device_id)]
+
+
+def check_device(device: Device | None) -> None:
+    if device is not None and not isinstance(device, Device):
+        raise ValueError(f"Unsupported device {device!r}")
+
+    if device is not None and device not in ALL_DEVICES:
+        raise ValueError(f"Unsupported device {device!r}")
+
+
+# Helpers for device-specific dtype support
+
+def get_default_dtypes(device: Device | None = None) -> dict[str, DType]:
+    if device == NO_FLOAT64_DEVICE:
+        # mimic an MPS device which does not have float64 at all
+        return {
+            "real floating": float32,
+            "complex floating": complex64,
+            "integral": int64,
+            "indexing": int64,
+        }
+    elif device == Device('device2'):
+        # mimic a torch CPU device: support float64 but default to float32
+        return {
+            "real floating": float32,
+            "complex floating": complex64,
+            "integral": int64,
+            "indexing": int64,
+        }
+    else:
+        return {
+            "real floating": float64,
+            "complex floating": complex128,
+            "integral": int64,
+            "indexing": int64,
+        }
+
+
+def device_supports_dtype(device: Device | None, dtype: DType |None) -> bool:
+    """True if `device` supports `dtype`, False otherwise."""
+    # Device("no_float64") supports all dtypes except float64 and complex128
+    if device == NO_FLOAT64_DEVICE:
+        return dtype not in (float64, complex128)
+
+    # All other devices support all dtypes
+    return True
+
+
+def _map_supported(dtypes: list[DType], device: Device) -> dict[str, DType]:
+    return {
+        dt._canonical_name: dt
+        for dt in dtypes
+        if device_supports_dtype(device, dt)
+    }
+
+
+# _info.dtypes() maps "kind" -> dict of {name: dtype}
+# Note that "kinds" differ from "categories" above, per the spec.
+
+_kind_to_dtypes = {
+    None: _all_dtypes,
+    "bool": _boolean_dtypes,
+    "signed integer": _signed_integer_dtypes,
+    "unsigned integer": _unsigned_integer_dtypes,
+    "integral": _integer_dtypes,
+    "real floating": _real_floating_dtypes,
+    "complex floating": _complex_floating_dtypes,
+    "numeric": _numeric_dtypes
+}
+

array_api_strict/_dtypes.py

@@ -11,9 +11,11 @@ import numpy.typing as npt
 
 class DType:
     _np_dtype: Final[np.dtype[Any]]
-    __slots__ = ("_np_dtype", "__weakref__")
+    _canonical_name: Final[Any]
+    __slots__ = ("_np_dtype", "_canonical_name", "__weakref__")
 
     def __init__(self, np_dtype: npt.DTypeLike):
+        self._canonical_name = np_dtype
         self._np_dtype = np.dtype(np_dtype)
 
     def __repr__(self) -> str: