scipy 1.16.2__cp314-cp314t-win_arm64.whl

scipy/_lib/array_api_extra/_lib/_lazy.py

@@ -0,0 +1,357 @@
+"""Public API Functions."""
+
+from __future__ import annotations
+
+import math
+from collections.abc import Callable, Sequence
+from functools import partial, wraps
+from types import ModuleType
+from typing import TYPE_CHECKING, Any, ParamSpec, TypeAlias, cast, overload
+
+from ._funcs import broadcast_shapes
+from ._utils import _compat
+from ._utils._compat import (
+    array_namespace,
+    is_dask_namespace,
+    is_jax_namespace,
+)
+from ._utils._helpers import is_python_scalar
+from ._utils._typing import Array, DType
+
+if TYPE_CHECKING:  # pragma: no cover
+    import numpy as np
+    from numpy.typing import ArrayLike
+
+    NumPyObject: TypeAlias = np.ndarray[Any, Any] | np.generic  # type: ignore[explicit-any]
+else:
+    # Sphinx hack
+    NumPyObject = Any
+
+P = ParamSpec("P")
+
+
+@overload
+def lazy_apply(  # type: ignore[decorated-any, valid-type]
+    func: Callable[P, Array | ArrayLike],
+    *args: Array | complex | None,
+    shape: tuple[int | None, ...] | None = None,
+    dtype: DType | None = None,
+    as_numpy: bool = False,
+    xp: ModuleType | None = None,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
+) -> Array: ...  # numpydoc ignore=GL08
+
+
+@overload
+def lazy_apply(  # type: ignore[decorated-any, valid-type]
+    func: Callable[P, Sequence[Array | ArrayLike]],
+    *args: Array | complex | None,
+    shape: Sequence[tuple[int | None, ...]],
+    dtype: Sequence[DType] | None = None,
+    as_numpy: bool = False,
+    xp: ModuleType | None = None,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
+) -> tuple[Array, ...]: ...  # numpydoc ignore=GL08
+
+
+def lazy_apply(  # type: ignore[valid-type]  # numpydoc ignore=GL07,SA04
+    func: Callable[P, Array | ArrayLike | Sequence[Array | ArrayLike]],
+    *args: Array | complex | None,
+    shape: tuple[int | None, ...] | Sequence[tuple[int | None, ...]] | None = None,
+    dtype: DType | Sequence[DType] | None = None,
+    as_numpy: bool = False,
+    xp: ModuleType | None = None,
+    **kwargs: P.kwargs,  # pyright: ignore[reportGeneralTypeIssues]
+) -> Array | tuple[Array, ...]:
+    """
+    Lazily apply an eager function.
+
+    If the backend of the input arrays is lazy, e.g. Dask or jitted JAX, the execution
+    of the function is delayed until the graph is materialized; if it's eager, the
+    function is executed immediately.
+
+    Parameters
+    ----------
+    func : callable
+        The function to apply.
+
+        It must accept one or more array API compliant arrays as positional arguments.
+        If `as_numpy=True`, inputs are converted to NumPy before they are passed to
+        `func`.
+        It must return either a single array-like or a sequence of array-likes.
+
+        `func` must be a pure function, i.e. without side effects, as depending on the
+        backend it may be executed more than once or never.
+    *args : Array | int | float | complex | bool | None
+        One or more Array API compliant arrays, Python scalars, or None's.
+
+        If `as_numpy=True`, you need to be able to apply :func:`numpy.asarray` to
+        non-None args to convert them to NumPy; read notes below about specific
+        backends.
+    shape : tuple[int | None, ...] | Sequence[tuple[int | None, ...]], optional
+        Output shape or sequence of output shapes, one for each output of `func`.
+        Default: assume single output and broadcast shapes of the input arrays.
+    dtype : DType | Sequence[DType], optional
+        Output dtype or sequence of output dtypes, one for each output of `func`.
+        dtype(s) must belong to the same array namespace as the input arrays.
+        Default: infer the result type(s) from the input arrays.
+    as_numpy : bool, optional
+        If True, convert the input arrays to NumPy before passing them to `func`.
+        This is particularly useful to make NumPy-only functions, e.g. written in Cython
+       or Numba, work transparently with array API-compliant arrays.
+        Default: False.
+    xp : array_namespace, optional
+        The standard-compatible namespace for `args`. Default: infer.
+    **kwargs : Any, optional
+        Additional keyword arguments to pass verbatim to `func`.
+        They cannot contain Array objects.
+
+    Returns
+    -------
+    Array | tuple[Array, ...]
+        The result(s) of `func` applied to the input arrays, wrapped in the same
+        array namespace as the inputs.
+        If shape is omitted or a single `tuple[int | None, ...]`, return a single array.
+        Otherwise, return a tuple of arrays.
+
+    Notes
+    -----
+    JAX
+        This allows applying eager functions to jitted JAX arrays, which are lazy.
+        The function won't be applied until the JAX array is materialized.
+        When running inside ``jax.jit``, `shape` must be fully known, i.e. it cannot
+        contain any `None` elements.
+
+        .. warning::
+
+            `func` must never raise inside ``jax.jit``, as the resulting behavior is
+            undefined.
+
+        Using this with `as_numpy=False` is particularly useful to apply non-jittable
+        JAX functions to arrays on GPU devices.
+        If ``as_numpy=True``, the :doc:`jax:transfer_guard` may prevent arrays on a GPU
+        device from being transferred back to CPU. This is treated as an implicit
+        transfer.
+
+    PyTorch, CuPy
+        If ``as_numpy=True``, these backends raise by default if you attempt to convert
+        arrays on a GPU device to NumPy.
+
+    Sparse
+        If ``as_numpy=True``, by default sparse prevents implicit densification through
+        :func:`numpy.asarray`. `This safety mechanism can be disabled
+        <https://sparse.pydata.org/en/stable/operations.html#package-configuration>`_.
+
+    Dask
+        This allows applying eager functions to Dask arrays.
+        The Dask graph won't be computed until the user calls ``compute()`` or
+        ``persist()`` down the line.
+
+        The function name will be prominently visible on the user-facing Dask
+        dashboard and on Prometheus metrics, so it is recommended for it to be
+        meaningful.
+
+        `lazy_apply` doesn't know if `func` reduces along any axes; also, shape
+        changes are non-trivial in chunked Dask arrays. For these reasons, all inputs
+        will be rechunked into a single chunk.
+
+        .. warning::
+
+           The whole operation needs to fit in memory all at once on a single worker.
+
+        The outputs will also be returned as a single chunk and you should consider
+        rechunking them into smaller chunks afterwards.
+
+        If you want to distribute the calculation across multiple workers, you
+        should use :func:`dask.array.map_blocks`, :func:`dask.array.map_overlap`,
+        :func:`dask.array.blockwise`, or a native Dask wrapper instead of
+        `lazy_apply`.
+
+    Dask wrapping around other backends
+        If ``as_numpy=False``, `func` will receive in input eager arrays of the meta
+        namespace, as defined by the ``._meta`` attribute of the input Dask arrays.
+        The outputs of `func` will be wrapped by the meta namespace, and then wrapped
+        again by Dask.
+
+    Raises
+    ------
+    ValueError
+        When ``xp=jax.numpy``, the output `shape` is unknown (it contains ``None`` on
+        one or more axes) and this function was called inside ``jax.jit``.
+    RuntimeError
+        When ``xp=sparse`` and auto-densification is disabled.
+    Exception (backend-specific)
+        When the backend disallows implicit device to host transfers and the input
+        arrays are on a non-CPU device, e.g. on GPU.
+
+    See Also
+    --------
+    jax.transfer_guard
+    jax.pure_callback
+    dask.array.map_blocks
+    dask.array.map_overlap
+    dask.array.blockwise
+    """
+    args_not_none = [arg for arg in args if arg is not None]
+    array_args = [arg for arg in args_not_none if not is_python_scalar(arg)]
+    if not array_args:
+        msg = "Must have at least one argument array"
+        raise ValueError(msg)
+    if xp is None:
+        xp = array_namespace(*args)
+
+    # Normalize and validate shape and dtype
+    shapes: list[tuple[int | None, ...]]
+    dtypes: list[DType]
+    multi_output = False
+
+    if shape is None:
+        shapes = [broadcast_shapes(*(arg.shape for arg in array_args))]
+    elif all(isinstance(s, int | None) for s in shape):
+        # Do not test for shape to be a tuple
+        # https://github.com/data-apis/array-api/issues/891#issuecomment-2637430522
+        shapes = [cast(tuple[int | None, ...], shape)]
+    else:
+        shapes = list(shape)  # type: ignore[arg-type]  # pyright: ignore[reportAssignmentType]
+        multi_output = True
+
+    if dtype is None:
+        dtypes = [xp.result_type(*args_not_none)] * len(shapes)
+    elif multi_output:
+        if not isinstance(dtype, Sequence):
+            msg = "Got multiple shapes but only one dtype"
+            raise ValueError(msg)
+        dtypes = list(dtype)  # pyright: ignore[reportUnknownArgumentType]
+    else:
+        if isinstance(dtype, Sequence):
+            msg = "Got single shape but multiple dtypes"
+            raise ValueError(msg)
+
+        dtypes = [dtype]
+
+    if len(shapes) != len(dtypes):
+        msg = f"Got {len(shapes)} shapes and {len(dtypes)} dtypes"
+        raise ValueError(msg)
+    del shape
+    del dtype
+    # End of shape and dtype parsing
+
+    # Backend-specific branches
+    if is_dask_namespace(xp):
+        import dask
+
+        metas: list[Array] = [arg._meta for arg in array_args]  # pylint: disable=protected-access    # pyright: ignore[reportAttributeAccessIssue]
+        meta_xp = array_namespace(*metas)
+
+        wrapped = dask.delayed(  # type: ignore[attr-defined]  # pyright: ignore[reportPrivateImportUsage]
+            _lazy_apply_wrapper(func, as_numpy, multi_output, meta_xp),
+            pure=True,
+        )
+        # This finalizes each arg, which is the same as arg.rechunk(-1).
+        # Please read docstring above for why we're not using
+        # dask.array.map_blocks or dask.array.blockwise!
+        delayed_out = wrapped(*args, **kwargs)
+
+        out = tuple(
+            xp.from_delayed(
+                delayed_out[i],  # pyright: ignore[reportIndexIssue]
+                # Dask's unknown shapes diverge from the Array API specification
+                shape=tuple(math.nan if s is None else s for s in shape),
+                dtype=dtype,
+                meta=metas[0],
+            )
+            for i, (shape, dtype) in enumerate(zip(shapes, dtypes, strict=True))
+        )
+
+    elif is_jax_namespace(xp) and _is_jax_jit_enabled(xp):
+        # Delay calling func with jax.pure_callback, which will forward to func eager
+        # JAX arrays. Do not use jax.pure_callback when running outside of the JIT,
+        # as it does not support raising exceptions:
+        # https://github.com/jax-ml/jax/issues/26102
+        import jax
+
+        if any(None in shape for shape in shapes):
+            msg = "Output shape must be fully known when running inside jax.jit"
+            raise ValueError(msg)
+
+        # Shield kwargs from being coerced into JAX arrays.
+        # jax.pure_callback calls jax.jit under the hood, but without the chance of
+        # passing static_argnames / static_argnums.
+        wrapped = _lazy_apply_wrapper(
+            partial(func, **kwargs), as_numpy, multi_output, xp
+        )
+
+        # suppress unused-ignore to run mypy in -e lint as well as -e dev
+        out = cast(  # type: ignore[bad-cast,unused-ignore]
+            tuple[Array, ...],
+            jax.pure_callback(
+                wrapped,
+                tuple(
+                    jax.ShapeDtypeStruct(shape, dtype)  # pyright: ignore[reportUnknownArgumentType]
+                    for shape, dtype in zip(shapes, dtypes, strict=True)
+                ),
+                *args,
+            ),
+        )
+
+    else:
+        # Eager backends, including non-jitted JAX
+        wrapped = _lazy_apply_wrapper(func, as_numpy, multi_output, xp)
+        out = wrapped(*args, **kwargs)
+
+    return out if multi_output else out[0]
+
+
+def _is_jax_jit_enabled(xp: ModuleType) -> bool:  # numpydoc ignore=PR01,RT01
+    """Return True if this function is being called inside ``jax.jit``."""
+    import jax  # pylint: disable=import-outside-toplevel
+
+    x = xp.asarray(False)
+    try:
+        return bool(x)
+    except jax.errors.TracerBoolConversionError:
+        return True
+
+
+def _lazy_apply_wrapper(  # type: ignore[explicit-any]  # numpydoc ignore=PR01,RT01
+    func: Callable[..., Array | ArrayLike | Sequence[Array | ArrayLike]],
+    as_numpy: bool,
+    multi_output: bool,
+    xp: ModuleType,
+) -> Callable[..., tuple[Array, ...]]:
+    """
+    Helper of `lazy_apply`.
+
+    Given a function that accepts one or more arrays as positional arguments and returns
+    a single array-like or a sequence of array-likes, return a function that accepts the
+    same number of Array API arrays and always returns a tuple of Array API array.
+
+    Any keyword arguments are passed through verbatim to the wrapped function.
+    """
+
+    # On Dask, @wraps causes the graph key to contain the wrapped function's name
+    @wraps(func)
+    def wrapper(  # type: ignore[decorated-any,explicit-any]
+        *args: Array | complex | None, **kwargs: Any
+    ) -> tuple[Array, ...]:  # numpydoc ignore=GL08
+        args_list = []
+        device = None
+        for arg in args:
+            if arg is not None and not is_python_scalar(arg):
+                if device is None:
+                    device = _compat.device(arg)
+                if as_numpy:
+                    import numpy as np
+
+                    arg = cast(Array, np.asarray(arg))  # type: ignore[bad-cast]  # noqa: PLW2901
+            args_list.append(arg)
+        assert device is not None
+
+        out = func(*args_list, **kwargs)
+
+        if multi_output:
+            assert isinstance(out, Sequence)
+            return tuple(xp.asarray(o, device=device) for o in out)
+        return (xp.asarray(out, device=device),)
+
+    return wrapper

scipy/_lib/array_api_extra/_lib/_testing.py

@@ -0,0 +1,278 @@
+"""
+Testing utilities.
+
+Note that this is private API; don't expect it to be stable.
+See also ..testing for public testing utilities.
+"""
+
+from __future__ import annotations
+
+import math
+from types import ModuleType
+from typing import Any, cast
+
+import numpy as np
+import pytest
+
+from ._utils._compat import (
+    array_namespace,
+    is_array_api_strict_namespace,
+    is_cupy_namespace,
+    is_dask_namespace,
+    is_jax_namespace,
+    is_numpy_namespace,
+    is_pydata_sparse_namespace,
+    is_torch_namespace,
+    to_device,
+)
+from ._utils._typing import Array, Device
+
+__all__ = ["as_numpy_array", "xp_assert_close", "xp_assert_equal", "xp_assert_less"]
+
+
+def _check_ns_shape_dtype(
+    actual: Array,
+    desired: Array,
+    check_dtype: bool,
+    check_shape: bool,
+    check_scalar: bool,
+) -> ModuleType:  # numpydoc ignore=RT03
+    """
+    Assert that namespace, shape and dtype of the two arrays match.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    check_dtype, check_shape : bool, default: True
+        Whether to check agreement between actual and desired dtypes and shapes
+    check_scalar : bool, default: False
+        NumPy only: whether to check agreement between actual and desired types -
+        0d array vs scalar.
+
+    Returns
+    -------
+    Arrays namespace.
+    """
+    actual_xp = array_namespace(actual)  # Raises on scalars and lists
+    desired_xp = array_namespace(desired)
+
+    msg = f"namespaces do not match: {actual_xp} != f{desired_xp}"
+    assert actual_xp == desired_xp, msg
+
+    if check_shape:
+        actual_shape = actual.shape
+        desired_shape = desired.shape
+        if is_dask_namespace(desired_xp):
+            # Dask uses nan instead of None for unknown shapes
+            if any(math.isnan(i) for i in cast(tuple[float, ...], actual_shape)):
+                actual_shape = actual.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+            if any(math.isnan(i) for i in cast(tuple[float, ...], desired_shape)):
+                desired_shape = desired.compute().shape  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+
+        msg = f"shapes do not match: {actual_shape} != f{desired_shape}"
+        assert actual_shape == desired_shape, msg
+
+    if check_dtype:
+        msg = f"dtypes do not match: {actual.dtype} != {desired.dtype}"
+        assert actual.dtype == desired.dtype, msg
+
+    if is_numpy_namespace(actual_xp) and check_scalar:
+        # only NumPy distinguishes between scalars and arrays; we do if check_scalar.
+        _msg = (
+            "array-ness does not match:\n Actual: "
+            f"{type(actual)}\n Desired: {type(desired)}"
+        )
+        assert np.isscalar(actual) == np.isscalar(desired), _msg
+
+    return desired_xp
+
+
+def as_numpy_array(array: Array, *, xp: ModuleType) -> np.typing.NDArray[Any]:  # type: ignore[explicit-any]
+    """
+    Convert array to NumPy, bypassing GPU-CPU transfer guards and densification guards.
+    """
+    if is_cupy_namespace(xp):
+        return xp.asnumpy(array)
+    if is_pydata_sparse_namespace(xp):
+        return array.todense()  # type: ignore[attr-defined]  # pyright: ignore[reportAttributeAccessIssue]
+
+    if is_torch_namespace(xp):
+        array = to_device(array, "cpu")
+    if is_array_api_strict_namespace(xp):
+        cpu: Device = xp.Device("CPU_DEVICE")
+        array = to_device(array, cpu)
+    if is_jax_namespace(xp):
+        import jax
+
+        # Note: only needed if the transfer guard is enabled
+        cpu = cast(Device, jax.devices("cpu")[0])
+        array = to_device(array, cpu)
+
+    return np.asarray(array)
+
+
+def xp_assert_equal(
+    actual: Array,
+    desired: Array,
+    *,
+    err_msg: str = "",
+    check_dtype: bool = True,
+    check_shape: bool = True,
+    check_scalar: bool = False,
+) -> None:
+    """
+    Array-API compatible version of `np.testing.assert_array_equal`.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    err_msg : str, optional
+        Error message to display on failure.
+    check_dtype, check_shape : bool, default: True
+        Whether to check agreement between actual and desired dtypes and shapes
+    check_scalar : bool, default: False
+        NumPy only: whether to check agreement between actual and desired types -
+        0d array vs scalar.
+
+    See Also
+    --------
+    xp_assert_close : Similar function for inexact equality checks.
+    numpy.testing.assert_array_equal : Similar function for NumPy arrays.
+    """
+    xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
+    actual_np = as_numpy_array(actual, xp=xp)
+    desired_np = as_numpy_array(desired, xp=xp)
+    np.testing.assert_array_equal(actual_np, desired_np, err_msg=err_msg)
+
+
+def xp_assert_less(
+    x: Array,
+    y: Array,
+    *,
+    err_msg: str = "",
+    check_dtype: bool = True,
+    check_shape: bool = True,
+    check_scalar: bool = False,
+) -> None:
+    """
+    Array-API compatible version of `np.testing.assert_array_less`.
+
+    Parameters
+    ----------
+    x, y : Array
+        The arrays to compare according to ``x < y`` (elementwise).
+    err_msg : str, optional
+        Error message to display on failure.
+    check_dtype, check_shape : bool, default: True
+        Whether to check agreement between actual and desired dtypes and shapes
+    check_scalar : bool, default: False
+        NumPy only: whether to check agreement between actual and desired types -
+        0d array vs scalar.
+
+    See Also
+    --------
+    xp_assert_close : Similar function for inexact equality checks.
+    numpy.testing.assert_array_equal : Similar function for NumPy arrays.
+    """
+    xp = _check_ns_shape_dtype(x, y, check_dtype, check_shape, check_scalar)
+    x_np = as_numpy_array(x, xp=xp)
+    y_np = as_numpy_array(y, xp=xp)
+    np.testing.assert_array_less(x_np, y_np, err_msg=err_msg)
+
+
+def xp_assert_close(
+    actual: Array,
+    desired: Array,
+    *,
+    rtol: float | None = None,
+    atol: float = 0,
+    err_msg: str = "",
+    check_dtype: bool = True,
+    check_shape: bool = True,
+    check_scalar: bool = False,
+) -> None:
+    """
+    Array-API compatible version of `np.testing.assert_allclose`.
+
+    Parameters
+    ----------
+    actual : Array
+        The array produced by the tested function.
+    desired : Array
+        The expected array (typically hardcoded).
+    rtol : float, optional
+        Relative tolerance. Default: dtype-dependent.
+    atol : float, optional
+        Absolute tolerance. Default: 0.
+    err_msg : str, optional
+        Error message to display on failure.
+    check_dtype, check_shape : bool, default: True
+        Whether to check agreement between actual and desired dtypes and shapes
+    check_scalar : bool, default: False
+        NumPy only: whether to check agreement between actual and desired types -
+        0d array vs scalar.
+
+    See Also
+    --------
+    xp_assert_equal : Similar function for exact equality checks.
+    isclose : Public function for checking closeness.
+    numpy.testing.assert_allclose : Similar function for NumPy arrays.
+
+    Notes
+    -----
+    The default `atol` and `rtol` differ from `xp.all(xpx.isclose(a, b))`.
+    """
+    xp = _check_ns_shape_dtype(actual, desired, check_dtype, check_shape, check_scalar)
+
+    if rtol is None:
+        if xp.isdtype(actual.dtype, ("real floating", "complex floating")):
+            # multiplier of 4 is used as for `np.float64` this puts the default `rtol`
+            # roughly half way between sqrt(eps) and the default for
+            # `numpy.testing.assert_allclose`, 1e-7
+            rtol = xp.finfo(actual.dtype).eps ** 0.5 * 4
+        else:
+            rtol = 1e-7
+
+    actual_np = as_numpy_array(actual, xp=xp)
+    desired_np = as_numpy_array(desired, xp=xp)
+    np.testing.assert_allclose(  # pyright: ignore[reportCallIssue]
+        actual_np,
+        desired_np,
+        rtol=rtol,  # pyright: ignore[reportArgumentType]
+        atol=atol,
+        err_msg=err_msg,
+    )
+
+
+def xfail(
+    request: pytest.FixtureRequest, *, reason: str, strict: bool | None = None
+) -> None:
+    """
+    XFAIL the currently running test.
+
+    Unlike ``pytest.xfail``, allow rest of test to execute instead of immediately
+    halting it, so that it may result in a XPASS.
+    xref https://github.com/pandas-dev/pandas/issues/38902
+
+    Parameters
+    ----------
+    request : pytest.FixtureRequest
+        ``request`` argument of the test function.
+    reason : str
+        Reason for the expected failure.
+    strict: bool, optional
+        If True, the test will be marked as failed if it passes.
+        If False, the test will be marked as passed if it fails.
+        Default: ``xfail_strict`` value in ``pyproject.toml``, or False if absent.
+    """
+    if strict is not None:
+        marker = pytest.mark.xfail(reason=reason, strict=strict)
+    else:
+        marker = pytest.mark.xfail(reason=reason)
+    request.node.add_marker(marker)

scipy/_lib/array_api_extra/_lib/_utils/__init__.py

@@ -0,0 +1 @@
+"""Modules housing private utility functions."""