scipy 1.15.3__cp313-cp313-musllinux_1_2_aarch64.whl → 1.16.0rc2__cp313-cp313-musllinux_1_2_aarch64.whl

scipy/_lib/array_api_compat/cupy/_typing.py

@@ -1,46 +1,31 @@
 from __future__ import annotations
 
-__all__ = [
-    "ndarray",
-    "Device",
-    "Dtype",
-]
+__all__ = ["Array", "DType", "Device"]
+_all_ignore = ["cp"]
 
-import sys
-from typing import (
-    Union,
-    TYPE_CHECKING,
-)
-
-from cupy import (
-    ndarray,
-    dtype,
-    int8,
-    int16,
-    int32,
-    int64,
-    uint8,
-    uint16,
-    uint32,
-    uint64,
-    float32,
-    float64,
-)
+from typing import TYPE_CHECKING
 
+import cupy as cp
+from cupy import ndarray as Array
 from cupy.cuda.device import Device
 
-if TYPE_CHECKING or sys.version_info >= (3, 9):
-    Dtype = dtype[Union[
-        int8,
-        int16,
-        int32,
-        int64,
-        uint8,
-        uint16,
-        uint32,
-        uint64,
-        float32,
-        float64,
-    ]]
+if TYPE_CHECKING:
+    # NumPy 1.x on Python 3.10 fails to parse np.dtype[]
+    DType = cp.dtype[
+        cp.intp
+        | cp.int8
+        | cp.int16
+        | cp.int32
+        | cp.int64
+        | cp.uint8
+        | cp.uint16
+        | cp.uint32
+        | cp.uint64
+        | cp.float32
+        | cp.float64
+        | cp.complex64
+        | cp.complex128
+        | cp.bool_
+    ]
 else:
-    Dtype = dtype
+    DType = cp.dtype

scipy/_lib/array_api_compat/dask/array/__init__.py

@@ -1,9 +1,12 @@
-from dask.array import * # noqa: F403
+from typing import Final
+
+from dask.array import *  # noqa: F403
 
 # These imports may overwrite names from the import * above.
-from ._aliases import * # noqa: F403
+from ._aliases import *  # noqa: F403
 
-__array_api_version__ = '2023.12'
+__array_api_version__: Final = "2024.12"
 
+# See the comment in the numpy __init__.py
 __import__(__package__ + '.linalg')
 __import__(__package__ + '.fft')

scipy/_lib/array_api_compat/dask/array/_aliases.py

@@ -1,66 +1,102 @@
-from __future__ import annotations
+# pyright: reportPrivateUsage=false
+# pyright: reportUnknownArgumentType=false
+# pyright: reportUnknownMemberType=false
+# pyright: reportUnknownVariableType=false
 
-from ...common import _aliases
-from ...common._helpers import _check_device
+from __future__ import annotations
 
-from ..._internal import get_xp
+from builtins import bool as py_bool
+from collections.abc import Callable
+from typing import TYPE_CHECKING, Any
 
-from ._info import __array_namespace_info__
+if TYPE_CHECKING:
+    from typing_extensions import TypeIs
 
+import dask.array as da
 import numpy as np
+from numpy import bool_ as bool
 from numpy import (
-    # Dtypes
-    iinfo,
-    finfo,
-    bool_ as bool,
+    can_cast,
+    complex64,
+    complex128,
     float32,
     float64,
     int8,
     int16,
     int32,
     int64,
+    result_type,
     uint8,
     uint16,
     uint32,
     uint64,
-    complex64,
-    complex128,
-    can_cast,
-    result_type,
 )
 
-from typing import TYPE_CHECKING
-if TYPE_CHECKING:
-    from typing import Optional, Union
-
-    from ...common._typing import Device, Dtype, Array, NestedSequence, SupportsBufferProtocol
-
-import dask.array as da
+from ..._internal import get_xp
+from ...common import _aliases, _helpers, array_namespace
+from ...common._typing import (
+    Array,
+    Device,
+    DType,
+    NestedSequence,
+    SupportsBufferProtocol,
+)
+from ._info import __array_namespace_info__
 
 isdtype = get_xp(np)(_aliases.isdtype)
 unstack = get_xp(da)(_aliases.unstack)
-astype = _aliases.astype
+
+
+# da.astype doesn't respect copy=True
+def astype(
+    x: Array,
+    dtype: DType,
+    /,
+    *,
+    copy: py_bool = True,
+    device: Device | None = None,
+) -> Array:
+    """
+    Array API compatibility wrapper for astype().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+    # TODO: respect device keyword?
+    _helpers._check_device(da, device)
+
+    if not copy and dtype == x.dtype:
+        return x
+    x = x.astype(dtype)
+    return x.copy() if copy else x
+
 
 # Common aliases
 
+
 # This arange func is modified from the common one to
 # not pass stop/step as keyword arguments, which will cause
 # an error with dask
-
-# TODO: delete the xp stuff, it shouldn't be necessary
-def _dask_arange(
-    start: Union[int, float],
+def arange(
+    start: float,
     /,
-    stop: Optional[Union[int, float]] = None,
-    step: Union[int, float] = 1,
+    stop: float | None = None,
+    step: float = 1,
     *,
-    xp,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
-    **kwargs,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    **kwargs: object,
 ) -> Array:
-    _check_device(xp, device)
-    args = [start]
+    """
+    Array API compatibility wrapper for arange().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+    # TODO: respect device keyword?
+    _helpers._check_device(da, device)
+
+    args: list[Any] = [start]
     if stop is not None:
         args.append(stop)
     else:
@@ -68,13 +104,12 @@ def _dask_arange(
         # prepend the default value for start which is 0
         args.insert(0, 0)
     args.append(step)
-    return xp.arange(*args, dtype=dtype, **kwargs)
 
-arange = get_xp(da)(_dask_arange)
-eye = get_xp(da)(_aliases.eye)
+    return da.arange(*args, dtype=dtype, **kwargs)
+
 
-linspace = get_xp(da)(_aliases.linspace)
 eye = get_xp(da)(_aliases.eye)
+linspace = get_xp(da)(_aliases.linspace)
 UniqueAllResult = get_xp(da)(_aliases.UniqueAllResult)
 UniqueCountsResult = get_xp(da)(_aliases.UniqueCountsResult)
 UniqueInverseResult = get_xp(da)(_aliases.UniqueInverseResult)
@@ -86,6 +121,7 @@ permute_dims = get_xp(da)(_aliases.permute_dims)
 std = get_xp(da)(_aliases.std)
 var = get_xp(da)(_aliases.var)
 cumulative_sum = get_xp(da)(_aliases.cumulative_sum)
+cumulative_prod = get_xp(da)(_aliases.cumulative_prod)
 empty = get_xp(da)(_aliases.empty)
 empty_like = get_xp(da)(_aliases.empty_like)
 full = get_xp(da)(_aliases.full)
@@ -97,7 +133,6 @@ zeros_like = get_xp(da)(_aliases.zeros_like)
 reshape = get_xp(da)(_aliases.reshape)
 matrix_transpose = get_xp(da)(_aliases.matrix_transpose)
 vecdot = get_xp(da)(_aliases.vecdot)
-
 nonzero = get_xp(da)(_aliases.nonzero)
 ceil = get_xp(np)(_aliases.ceil)
 floor = get_xp(np)(_aliases.floor)
@@ -105,23 +140,19 @@ trunc = get_xp(np)(_aliases.trunc)
 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.
 def asarray(
-    obj: Union[
-        Array,
-        bool,
-        int,
-        float,
-        NestedSequence[bool | int | float],
-        SupportsBufferProtocol,
-    ],
+    obj: complex | NestedSequence[complex] | Array | SupportsBufferProtocol,
     /,
     *,
-    dtype: Optional[Dtype] = None,
-    device: Optional[Device] = None,
-    copy: "Optional[Union[bool, np._CopyMode]]" = None,
-    **kwargs,
+    dtype: DType | None = None,
+    device: Device | None = None,
+    copy: py_bool | None = None,
+    **kwargs: object,
 ) -> Array:
     """
     Array API compatibility wrapper for asarray().
@@ -129,89 +160,217 @@ def asarray(
     See the corresponding documentation in the array library and/or the array API
     specification for more details.
     """
+    # TODO: respect device keyword?
+    _helpers._check_device(da, device)
+
+    if isinstance(obj, da.Array):
+        if dtype is not None and dtype != obj.dtype:
+            if copy is False:
+                raise ValueError("Unable to avoid copy when changing dtype")
+            obj = obj.astype(dtype)
+        return obj.copy() if copy else obj  # pyright: ignore[reportAttributeAccessIssue]
+
     if copy is False:
-        # copy=False is not yet implemented in dask
-        raise NotImplementedError("copy=False is not yet implemented")
-    elif copy is True:
-        if isinstance(obj, da.Array) and dtype is None:
-            return obj.copy()
-        # Go through numpy, since dask copy is no-op by default
-        obj = np.array(obj, dtype=dtype, copy=True)
-        return da.array(obj, dtype=dtype)
-    else:
-        if not isinstance(obj, da.Array) or dtype is not None and obj.dtype != dtype:
-            obj = np.asarray(obj, dtype=dtype)
-            return da.from_array(obj)
-        return obj
-
-    return da.asarray(obj, dtype=dtype, **kwargs)
-
-from dask.array import (
-    # Element wise aliases
-    arccos as acos,
-    arccosh as acosh,
-    arcsin as asin,
-    arcsinh as asinh,
-    arctan as atan,
-    arctan2 as atan2,
-    arctanh as atanh,
-    left_shift as bitwise_left_shift,
-    right_shift as bitwise_right_shift,
-    invert as bitwise_invert,
-    power as pow,
-    # Other
-    concatenate as concat,
-)
+        raise ValueError(
+            "Unable to avoid copy when converting a non-dask object to dask"
+        )
+
+    # copy=None to be uniform across dask < 2024.12 and >= 2024.12
+    # see https://github.com/dask/dask/pull/11524/
+    obj = np.array(obj, dtype=dtype, copy=True)
+    return da.from_array(obj)
+
+
+# Element wise aliases
+from dask.array import arccos as acos
+from dask.array import arccosh as acosh
+from dask.array import arcsin as asin
+from dask.array import arcsinh as asinh
+from dask.array import arctan as atan
+from dask.array import arctan2 as atan2
+from dask.array import arctanh as atanh
+
+# Other
+from dask.array import concatenate as concat
+from dask.array import invert as bitwise_invert
+from dask.array import left_shift as bitwise_left_shift
+from dask.array import power as pow
+from dask.array import right_shift as bitwise_right_shift
+
 
 # dask.array.clip does not work unless all three arguments are provided.
 # Furthermore, the masking workaround in common._aliases.clip cannot work with
 # dask (meaning uint64 promoting to float64 is going to just be unfixed for
 # now).
-@get_xp(da)
 def clip(
     x: Array,
     /,
-    min: Optional[Union[int, float, Array]] = None,
-    max: Optional[Union[int, float, Array]] = None,
-    *,
-    xp,
+    min: float | Array | None = None,
+    max: float | Array | None = None,
 ) -> Array:
-    def _isscalar(a):
-        return isinstance(a, (int, float, type(None)))
+    """
+    Array API compatibility wrapper for clip().
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+
+    def _isscalar(a: float | Array | None, /) -> TypeIs[float | None]:
+        return a is None or isinstance(a, (int, float))
+
     min_shape = () if _isscalar(min) else min.shape
     max_shape = () if _isscalar(max) else max.shape
 
     # TODO: This won't handle dask unknown shapes
-    import numpy as np
     result_shape = np.broadcast_shapes(x.shape, min_shape, max_shape)
 
     if min is not None:
-        min = xp.broadcast_to(xp.asarray(min), result_shape)
+        min = da.broadcast_to(da.asarray(min), result_shape)
     if max is not None:
-        max = xp.broadcast_to(xp.asarray(max), result_shape)
+        max = da.broadcast_to(da.asarray(max), result_shape)
 
     if min is None and max is None:
-        return xp.positive(x)
+        return da.positive(x)
 
     if min is None:
-        return astype(xp.minimum(x, max), x.dtype)
+        return astype(da.minimum(x, max), x.dtype)
     if max is None:
-        return astype(xp.maximum(x, min), x.dtype)
+        return astype(da.maximum(x, min), x.dtype)
+
+    return astype(da.minimum(da.maximum(x, min), max), x.dtype)
+
+
+def _ensure_single_chunk(x: Array, axis: int) -> tuple[Array, Callable[[Array], Array]]:
+    """
+    Make sure that Array is not broken into multiple chunks along axis.
+
+    Returns
+    -------
+    x : Array
+        The input Array with a single chunk along axis.
+    restore : Callable[Array, Array]
+        function to apply to the output to rechunk it back into reasonable chunks
+    """
+    if axis < 0:
+        axis += x.ndim
+    if x.numblocks[axis] < 2:
+        return x, lambda x: x
+
+    # Break chunks on other axes in an attempt to keep chunk size low
+    x = x.rechunk({i: -1 if i == axis else "auto" for i in range(x.ndim)})
+
+    # Rather than reconstructing the original chunks, which can be a
+    # very expensive affair, just break down oversized chunks without
+    # incurring in any transfers over the network.
+    # This has the downside of a risk of overchunking if the array is
+    # then used in operations against other arrays that match the
+    # original chunking pattern.
+    return x, lambda x: x.rechunk()
+
+
+def sort(
+    x: Array,
+    /,
+    *,
+    axis: int = -1,
+    descending: py_bool = False,
+    stable: py_bool = True,
+) -> Array:
+    """
+    Array API compatibility layer around the lack of sort() in Dask.
+
+    Warnings
+    --------
+    This function temporarily rechunks the array along `axis` to a single chunk.
+    This can be extremely inefficient and can lead to out-of-memory errors.
+
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+    """
+    x, restore = _ensure_single_chunk(x, axis)
+
+    meta_xp = array_namespace(x._meta)
+    x = da.map_blocks(
+        meta_xp.sort,
+        x,
+        axis=axis,
+        meta=x._meta,
+        dtype=x.dtype,
+        descending=descending,
+        stable=stable,
+    )
+
+    return restore(x)
+
+
+def argsort(
+    x: Array,
+    /,
+    *,
+    axis: int = -1,
+    descending: py_bool = False,
+    stable: py_bool = True,
+) -> Array:
+    """
+    Array API compatibility layer around the lack of argsort() in Dask.
 
-    return astype(xp.minimum(xp.maximum(x, min), max), x.dtype)
+    See the corresponding documentation in the array library and/or the array API
+    specification for more details.
+
+    Warnings
+    --------
+    This function temporarily rechunks the array along `axis` into a single chunk.
+    This can be extremely inefficient and can lead to out-of-memory errors.
+    """
+    x, restore = _ensure_single_chunk(x, axis)
+
+    meta_xp = array_namespace(x._meta)
+    dtype = meta_xp.argsort(x._meta).dtype
+    meta = meta_xp.astype(x._meta, dtype)
+    x = da.map_blocks(
+        meta_xp.argsort,
+        x,
+        axis=axis,
+        meta=meta,
+        dtype=dtype,
+        descending=descending,
+        stable=stable,
+    )
+
+    return restore(x)
+
+
+# dask.array.count_nonzero does not have keepdims
+def count_nonzero(
+    x: Array,
+    axis: int | None = None,
+    keepdims: py_bool = False,
+) -> Array:
+    result = da.count_nonzero(x, axis)
+    if keepdims:
+        if axis is None:
+            return da.reshape(result, [1] * x.ndim)
+        return da.expand_dims(result, axis)
+    return result
 
-# exclude these from all since dask.array has no sorting functions
-_da_unsupported = ['sort', 'argsort']
 
-_common_aliases = [alias for alias in _aliases.__all__ if alias not in _da_unsupported]
+__all__ = [
+    "__array_namespace_info__",
+    "count_nonzero",
+    "bool",
+    "int8", "int16", "int32", "int64",
+    "uint8", "uint16", "uint32", "uint64",
+    "float32", "float64",
+    "complex64", "complex128",
+    "asarray", "astype", "can_cast", "result_type",
+    "pow",
+    "concat",
+    "acos", "acosh", "asin", "asinh", "atan", "atan2", "atanh",
+    "bitwise_left_shift", "bitwise_right_shift", "bitwise_invert",
+]  # fmt: skip
+__all__ += _aliases.__all__
+_all_ignore = ["array_namespace", "get_xp", "da", "np"]
 
-__all__ = _common_aliases + ['__array_namespace_info__', 'asarray', 'acos',
-                    'acosh', 'asin', 'asinh', 'atan', 'atan2',
-                    'atanh', 'bitwise_left_shift', 'bitwise_invert',
-                    'bitwise_right_shift', 'concat', 'pow', 'iinfo', 'finfo', 'can_cast',
-                    'result_type', 'bool', 'float32', 'float64', 'int8', 'int16', 'int32', 'int64',
-                    'uint8', 'uint16', 'uint32', 'uint64',
-                    'complex64', 'complex128', 'iinfo', 'finfo',
-                    'can_cast', 'result_type']
 
-_all_ignore = ["get_xp", "da", "np"]
+def __dir__() -> list[str]:
+    return __all__