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

array_api_strict/_fft.py

@@ -4,7 +4,7 @@ from typing import TYPE_CHECKING
 
 if TYPE_CHECKING:
     from typing import Union, Optional, Literal
-    from ._typing import Device
+    from ._typing import Device, Dtype as DType
     from collections.abc import Sequence
 
 from ._dtypes import (
@@ -251,7 +251,14 @@ def ihfft(
     return res
 
 @requires_extension('fft')
-def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
+def fftfreq(
+    n: int,
+    /,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.fftfreq <numpy.fft.fftfreq>`.
 
@@ -259,10 +266,23 @@ def fftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Ar
     """
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.fftfreq(n, d=d), device=device)
+    if dtype and not dtype in _real_floating_dtypes:
+        raise ValueError(f"`dtype` must be a real floating-point type. Got {dtype=}.")
+
+    np_result = np.fft.fftfreq(n, d=d)
+    if dtype:
+        np_result = np_result.astype(dtype._np_dtype)
+    return Array._new(np_result, device=device)
 
 @requires_extension('fft')
-def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> Array:
+def rfftfreq(
+    n: int,
+    /,
+    *,
+    d: float = 1.0,
+    dtype: Optional[DType] = None,
+    device: Optional[Device] = None
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.fft.rfftfreq <numpy.fft.rfftfreq>`.
 
@@ -270,7 +290,13 @@ def rfftfreq(n: int, /, *, d: float = 1.0, device: Optional[Device] = None) -> A
     """
     if device is not None and device not in ALL_DEVICES:
         raise ValueError(f"Unsupported device {device!r}")
-    return Array._new(np.fft.rfftfreq(n, d=d), device=device)
+    if dtype and not dtype in _real_floating_dtypes:
+        raise ValueError(f"`dtype` must be a real floating-point type. Got {dtype=}.")
+
+    np_result = np.fft.rfftfreq(n, d=d)
+    if dtype:
+        np_result = np_result.astype(dtype._np_dtype)
+    return Array._new(np_result, device=device)
 
 @requires_extension('fft')
 def fftshift(x: Array, /, *, axes: Union[int, Sequence[int]] = None) -> Array:

array_api_strict/_flags.py

@@ -22,11 +22,12 @@ supported_versions = (
     "2021.12",
     "2022.12",
     "2023.12",
+    "2024.12"
 )
 
-draft_version = "2024.12"
+draft_version = "2025.12"
 
-API_VERSION = default_version = "2023.12"
+API_VERSION = default_version = "2024.12"
 
 BOOLEAN_INDEXING = True
 
@@ -169,12 +170,13 @@ def set_array_api_strict_flags(
                                          set(default_extensions))
 
 # We have to do this separately or it won't get added as the docstring
-set_array_api_strict_flags.__doc__ = set_array_api_strict_flags.__doc__.format(
-    supported_versions=supported_versions,
-    default_version=default_version,
-    default_extensions=default_extensions,
-    draft_version=draft_version,
-)
+if set_array_api_strict_flags.__doc__ is not None:
+    set_array_api_strict_flags.__doc__ = set_array_api_strict_flags.__doc__.format(
+        supported_versions=supported_versions,
+        default_version=default_version,
+        default_extensions=default_extensions,
+        draft_version=draft_version,
+    )
 
 def get_array_api_strict_flags():
     """

array_api_strict/_helpers.py

@@ -0,0 +1,37 @@
+"""Private helper routines.
+"""
+
+from ._flags import get_array_api_strict_flags
+from ._dtypes import _dtype_categories
+
+_py_scalars = (bool, int, float, complex)
+
+
+def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
+
+    flags = get_array_api_strict_flags()
+    if flags["api_version"] < "2024.12":
+        # scalars will fail at the call site
+        return x1, x2
+
+    _allowed_dtypes = _dtype_categories[dtype_category]
+
+    if isinstance(x1, _py_scalars):
+        if isinstance(x2, _py_scalars):
+            raise TypeError(f"Two scalars not allowed, got {type(x1) = } and {type(x2) =}")
+        # x2 must be an array
+        if x2.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x2.dtype}.")
+        x1 = x2._promote_scalar(x1)
+
+    elif isinstance(x2, _py_scalars):
+        # x1 must be an array
+        if x1.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x1.dtype}.")
+        x2 = x1._promote_scalar(x2)
+    else:
+        if x1.dtype not in _allowed_dtypes or x2.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. "
+                            f"Got {x1.dtype} and {x2.dtype}.")
+    return x1, x2
+

array_api_strict/_linear_algebra_functions.py

@@ -86,5 +86,5 @@ def vecdot(x1: Array, x2: Array, /, *, axis: int = -1) -> Array:
     x1_ = np.moveaxis(x1_, axis, -1)
     x2_ = np.moveaxis(x2_, axis, -1)
 
-    res = x1_[..., None, :] @ x2_[..., None]
+    res = np.conj(x1_[..., None, :]) @ x2_[..., None]
     return Array._new(res[..., 0, 0], device=x1.device)

array_api_strict/_manipulation_functions.py

@@ -153,6 +153,11 @@ def squeeze(x: Array, /, axis: Union[int, Tuple[int, ...]]) -> Array:
 
     See its docstring for more information.
     """
+    if axis is None:
+        raise ValueError(
+            "squeeze(..., axis=None is not supported. See "
+            "https://github.com/data-apis/array-api/pull/100 for a discussion."
+        )
     return Array._new(np.squeeze(x._array, axis=axis), device=x.device)
 
 

array_api_strict/_searching_functions.py

@@ -1,12 +1,12 @@
 from __future__ import annotations
 
 from ._array_object import Array
-from ._dtypes import _result_type, _real_numeric_dtypes
-from ._flags import requires_data_dependent_shapes, requires_api_version
+from ._dtypes import _result_type, _real_numeric_dtypes, bool as _bool
+from ._flags import requires_data_dependent_shapes, requires_api_version, get_array_api_strict_flags
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
-    from typing import Literal, Optional, Tuple
+    from typing import Literal, Optional, Tuple, Union
 
 import numpy as np
 
@@ -45,6 +45,24 @@ def nonzero(x: Array, /) -> Tuple[Array, ...]:
         raise ValueError("nonzero is not allowed on 0-dimensional arrays")
     return tuple(Array._new(i, device=x.device) for i in np.nonzero(x._array))
 
+
+@requires_api_version('2024.12')
+def count_nonzero(
+    x: Array,
+    /,
+    *,
+    axis: Optional[Union[int, Tuple[int, ...]]] = None,
+    keepdims: bool = False,
+) -> Array:
+    """
+    Array API compatible wrapper for :py:func:`np.count_nonzero <numpy.count_nonzero>`
+
+    See its docstring for more information.
+    """
+    arr = np.count_nonzero(x._array, axis=axis, keepdims=keepdims)
+    return Array._new(np.asarray(arr), device=x.device)
+
+
 @requires_api_version('2023.12')
 def searchsorted(
     x1: Array,
@@ -72,17 +90,38 @@ def searchsorted(
     # x1 must be 1-D, but NumPy already requires this.
     return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter), device=x1.device)
 
-def where(condition: Array, x1: Array, x2: Array, /) -> Array:
+def where(
+    condition: Array,
+    x1: bool | int | float | complex | Array,
+    x2: bool | int | float | complex | Array, /
+) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.where <numpy.where>`.
 
     See its docstring for more information.
     """
+    if get_array_api_strict_flags()['api_version'] > '2023.12':
+        num_scalars = 0
+
+        if isinstance(x1, (bool, float, complex, int)):
+            x1 = Array._new(np.asarray(x1), device=condition.device)
+            num_scalars += 1
+
+        if isinstance(x2, (bool, float, complex, int)):
+            x2 = Array._new(np.asarray(x2), device=condition.device)
+            num_scalars += 1
+
+        if num_scalars == 2:
+            raise ValueError("One of x1, x2 arguments must be an array.")
+
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
+    
+    if condition.dtype != _bool:
+        raise TypeError("`condition` must be have a boolean data type")
 
     if len({a.device for a in (condition, x1, x2)}) > 1:
-        raise ValueError("where inputs must all be on the same device")
+        raise ValueError("Inputs to `where` must all use the same device")
 
     x1, x2 = Array._normalize_two_args(x1, x2)
     return Array._new(np.where(condition._array, x1._array, x2._array), device=x1.device)

array_api_strict/_statistical_functions.py

@@ -3,12 +3,13 @@ from __future__ import annotations
 from ._dtypes import (
     _real_floating_dtypes,
     _real_numeric_dtypes,
+    _floating_dtypes,
     _numeric_dtypes,
 )
 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
+from ._creation_functions import zeros, ones
 from ._manipulation_functions import concat
 
 from typing import TYPE_CHECKING
@@ -30,7 +31,6 @@ def cumulative_sum(
 ) -> Array:
     if x.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in cumulative_sum")
-    dt = x.dtype if dtype is None else dtype
     if dtype is not None:
         dtype = dtype._np_dtype
 
@@ -43,9 +43,40 @@ def cumulative_sum(
     if include_initial:
         if axis < 0:
             axis += x.ndim
-        x = concat([zeros(x.shape[:axis] + (1,) + x.shape[axis + 1:], dtype=dt), x], axis=axis)
+        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)
 
+
+@requires_api_version('2024.12')
+def cumulative_prod(
+    x: Array,
+    /,
+    *,
+    axis: Optional[int] = None,
+    dtype: Optional[Dtype] = None,
+    include_initial: bool = False,
+) -> Array:
+    if x.dtype not in _numeric_dtypes:
+        raise TypeError("Only numeric dtypes are allowed in 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")
+        axis = 0
+
+    # np.cumprod does not support include_initial
+    if include_initial:
+        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)
+
+
 def max(
     x: Array,
     /,
@@ -65,8 +96,14 @@ def mean(
     axis: Optional[Union[int, Tuple[int, ...]]] = None,
     keepdims: bool = False,
 ) -> Array:
-    if x.dtype not in _real_floating_dtypes:
-        raise TypeError("Only real floating-point dtypes are allowed in mean")
+
+    if get_array_api_strict_flags()['api_version'] > '2023.12':
+        allowed_dtypes = _floating_dtypes
+    else:
+        allowed_dtypes = _real_floating_dtypes
+
+    if x.dtype not in allowed_dtypes:
+        raise TypeError("Only floating-point dtypes are allowed in mean")
     return Array._new(np.mean(x._array, axis=axis, keepdims=keepdims), device=x.device)
 
 

array_api_strict/_version.py

@@ -1,21 +1,21 @@
+# file generated by setuptools-scm
+# don't change, don't track in version control
 
-# This file was generated by 'versioneer.py' (0.29) from
-# revision-control system data, or from the parent directory name of an
-# unpacked source archive. Distribution tarballs contain a pre-generated copy
-# of this file.
+__all__ = ["__version__", "__version_tuple__", "version", "version_tuple"]
 
-import json
+TYPE_CHECKING = False
+if TYPE_CHECKING:
+    from typing import Tuple
+    from typing import Union
 
-version_json = '''
-{
- "date": "2024-11-11T15:30:59-0700",
- "dirty": false,
- "error": null,
- "full-revisionid": "fe4f34ad30ba59c1c981037dd181f6b4daf23b22",
- "version": "2.2"
-}
-'''  # END VERSION_JSON
+    VERSION_TUPLE = Tuple[Union[int, str], ...]
+else:
+    VERSION_TUPLE = object
 
+version: str
+__version__: str
+__version_tuple__: VERSION_TUPLE
+version_tuple: VERSION_TUPLE
 
-def get_versions():
-    return json.loads(version_json)
+__version__ = version = '2.3'
+__version_tuple__ = version_tuple = (2, 3)