array-api-strict 2.3__py3-none-any.whl → 2.4__py3-none-any.whl

array_api_strict/_statistical_functions.py

@@ -1,38 +1,34 @@
-from __future__ import annotations
+from typing import Any
 
+import numpy as np
+
+from ._array_object import Array
+from ._creation_functions import ones, zeros
 from ._dtypes import (
-    _real_floating_dtypes,
-    _real_numeric_dtypes,
+    DType,
     _floating_dtypes,
+    _np_dtype,
     _numeric_dtypes,
+    _real_floating_dtypes,
+    _real_numeric_dtypes,
+    complex64,
+    float32,
 )
-from ._array_object import Array
-from ._dtypes import float32, complex64
-from ._flags import requires_api_version, get_array_api_strict_flags
-from ._creation_functions import zeros, ones
+from ._flags import get_array_api_strict_flags, requires_api_version
 from ._manipulation_functions import concat
 
-from typing import TYPE_CHECKING
-
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
-    from ._typing import Dtype
-
-import numpy as np
 
 @requires_api_version('2023.12')
 def cumulative_sum(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in cumulative_sum")
-    if dtype is not None:
-        dtype = dtype._np_dtype
 
     # TODO: The standard is not clear about what should happen when x.ndim == 0.
     if axis is None:
@@ -44,7 +40,7 @@ def cumulative_sum(
         if axis < 0:
             axis += x.ndim
         x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumsum(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumsum(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 @requires_api_version('2024.12')
@@ -52,8 +48,8 @@ def cumulative_prod(
     x: Array,
     /,
     *,
-    axis: Optional[int] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | None = None,
+    dtype: DType | None = None,
     include_initial: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
@@ -61,9 +57,6 @@ def cumulative_prod(
     if x.ndim == 0:
         raise ValueError("Only ndim >= 1 arrays are allowed in cumulative_prod")
 
-    if dtype is not None:
-        dtype = dtype._np_dtype
-
     if axis is None:
         if x.ndim > 1:
             raise ValueError("axis must be specified in cumulative_prod for more than one dimension")
@@ -74,14 +67,14 @@ def cumulative_prod(
         if axis < 0:
             axis += x.ndim
         x = concat([ones(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=x.dtype), x], axis=axis)
-    return Array._new(np.cumprod(x._array, axis=axis, dtype=dtype), device=x.device)
+    return Array._new(np.cumprod(x._array, axis=axis, dtype=_np_dtype(dtype)), device=x.device)
 
 
 def max(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
@@ -93,14 +86,15 @@ def mean(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
 
-    if get_array_api_strict_flags()['api_version'] > '2023.12':
-        allowed_dtypes = _floating_dtypes
-    else:
-        allowed_dtypes = _real_floating_dtypes
+    allowed_dtypes = (
+        _floating_dtypes
+        if get_array_api_strict_flags()['api_version'] > '2023.12'
+        else _real_floating_dtypes
+    )
 
     if x.dtype not in allowed_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in mean")
@@ -111,7 +105,7 @@ def min(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _real_numeric_dtypes:
@@ -119,37 +113,43 @@ def min(
     return Array._new(np.min(x._array, axis=axis, keepdims=keepdims), device=x.device)
 
 
+def _np_dtype_sumprod(x: Array, dtype: DType | None) -> np.dtype[Any] | None:
+    """In versions prior to 2023.12, sum() and prod() upcast for all
+    dtypes when dtype=None. For 2023.12, the behavior is the same as in
+    NumPy (only upcast for integral dtypes).
+    """
+    if dtype is None and get_array_api_strict_flags()['api_version'] < '2023.12':
+        if x.dtype == float32:
+            return np.float64  # type: ignore[return-value]
+        elif x.dtype == complex64:
+            return np.complex128  # type: ignore[return-value]
+    return _np_dtype(dtype)
+
+
 def prod(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in prod")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.prod(x._array, dtype=dtype, axis=axis, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.prod(x._array, dtype=np_dtype, axis=axis, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def std(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here
@@ -162,33 +162,26 @@ def sum(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    dtype: Optional[Dtype] = None,
+    axis: int | tuple[int, ...] | None = None,
+    dtype: DType | None = None,
     keepdims: bool = False,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in sum")
 
-    if dtype is None:
-        # Note: In versions prior to 2023.12, sum() and prod() upcast for all
-        # dtypes when dtype=None. For 2023.12, the behavior is the same as in
-        # NumPy (only upcast for integral dtypes).
-        if get_array_api_strict_flags()['api_version'] < '2023.12':
-            if x.dtype == float32:
-                dtype = np.float64
-            elif x.dtype == complex64:
-                dtype = np.complex128
-    else:
-        dtype = dtype._np_dtype
-    return Array._new(np.sum(x._array, axis=axis, dtype=dtype, keepdims=keepdims), device=x.device)
+    np_dtype = _np_dtype_sumprod(x, dtype)
+    return Array._new(
+        np.sum(x._array, axis=axis, dtype=np_dtype, keepdims=keepdims),
+        device=x.device,
+    )
 
 
 def var(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
-    correction: Union[int, float] = 0.0,
+    axis: int | tuple[int, ...] | None = None,
+    correction: float = 0.0,
     keepdims: bool = False,
 ) -> Array:
     # Note: the keyword argument correction is different here

array_api_strict/_typing.py

@@ -8,41 +8,19 @@ valid for inputs that match the given type annotations.
 
 from __future__ import annotations
 
-__all__ = [
-    "Array",
-    "Device",
-    "Dtype",
-    "SupportsDLPack",
-    "SupportsBufferProtocol",
-    "PyCapsule",
-]
-
 import sys
+from typing import Any, Protocol, TypedDict, TypeVar
 
-from typing import (
-    Any,
-    TypedDict,
-    TypeVar,
-    Protocol,
-)
-
-from ._array_object import Array, _device
-from ._dtypes import _DType
-from ._info import __array_namespace_info__
+from ._dtypes import DType
 
 _T_co = TypeVar("_T_co", covariant=True)
 
+
 class NestedSequence(Protocol[_T_co]):
     def __getitem__(self, key: int, /) -> _T_co | NestedSequence[_T_co]: ...
     def __len__(self, /) -> int: ...
 
 
-Device = _device
-
-Dtype = _DType
-
-Info = __array_namespace_info__
-
 if sys.version_info >= (3, 12):
     from collections.abc import Buffer as SupportsBufferProtocol
 else:
@@ -50,40 +28,42 @@ else:
 
 PyCapsule = Any
 
+
 class SupportsDLPack(Protocol):
     def __dlpack__(self, /, *, stream: None = ...) -> PyCapsule: ...
 
+
 Capabilities = TypedDict(
-    "Capabilities", {"boolean indexing": bool, "data-dependent shapes": bool,
-                     "max dimensions": int}
+    "Capabilities",
+    {
+        "boolean indexing": bool,
+        "data-dependent shapes": bool,
+        "max dimensions": int,
+    },
 )
 
 DefaultDataTypes = TypedDict(
     "DefaultDataTypes",
     {
-        "real floating": Dtype,
-        "complex floating": Dtype,
-        "integral": Dtype,
-        "indexing": Dtype,
+        "real floating": DType,
+        "complex floating": DType,
+        "integral": DType,
+        "indexing": DType,
     },
 )
 
-DataTypes = TypedDict(
-    "DataTypes",
-    {
-        "bool": Dtype,
-        "float32": Dtype,
-        "float64": Dtype,
-        "complex64": Dtype,
-        "complex128": Dtype,
-        "int8": Dtype,
-        "int16": Dtype,
-        "int32": Dtype,
-        "int64": Dtype,
-        "uint8": Dtype,
-        "uint16": Dtype,
-        "uint32": Dtype,
-        "uint64": Dtype,
-    },
-    total=False,
-)
+
+class DataTypes(TypedDict, total=False):
+    bool: DType
+    float32: DType
+    float64: DType
+    complex64: DType
+    complex128: DType
+    int8: DType
+    int16: DType
+    int32: DType
+    int64: DType
+    uint8: DType
+    uint16: DType
+    uint32: DType
+    uint64: DType

array_api_strict/_utility_functions.py

@@ -1,21 +1,18 @@
-from __future__ import annotations
+from typing import Any
+
+import numpy as np
+import numpy.typing as npt
 
 from ._array_object import Array
-from ._flags import requires_api_version
 from ._dtypes import _numeric_dtypes
-
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Optional, Tuple, Union
-
-import numpy as np
+from ._flags import requires_api_version
 
 
 def all(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -30,7 +27,7 @@ def any(
     x: Array,
     /,
     *,
-    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    axis: int | tuple[int, ...] | None = None,
     keepdims: bool = False,
 ) -> Array:
     """
@@ -40,6 +37,7 @@ def any(
     """
     return Array._new(np.asarray(np.any(x._array, axis=axis, keepdims=keepdims)), device=x.device)
 
+
 @requires_api_version('2024.12')
 def diff(
     x: Array,
@@ -47,8 +45,8 @@ def diff(
     *,
     axis: int = -1,
     n: int = 1,
-    prepend: Optional[Array] = None,
-    append: Optional[Array] = None,
+    prepend: Array | None = None,
+    append: Array | None = None,
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in diff")
@@ -57,7 +55,7 @@ def diff(
     # currently specified.
 
     # NumPy does not support prepend=None or append=None
-    kwargs = dict(axis=axis, n=n)
+    kwargs: dict[str, int | npt.NDArray[Any]] = {"axis": axis, "n": n}
     if prepend is not None:
         if prepend.device != x.device:
             raise ValueError(f"Arrays from two different devices ({prepend.device} and {x.device}) can not be combined.")

array_api_strict/_version.py

@@ -17,5 +17,5 @@ __version__: str
 __version_tuple__: VERSION_TUPLE
 version_tuple: VERSION_TUPLE
 
-__version__ = version = '2.3'
-__version_tuple__ = version_tuple = (2, 3)
+__version__ = version = '2.4'
+__version_tuple__ = version_tuple = (2, 4)