Merge pull request numpy#26724 from mtsokol/array-api-2023

ENH: Add Array API 2023.12 version support

Author: mhvk

.github/workflows/linux.yml

@@ -231,7 +231,7 @@ jobs:
       uses: actions/checkout@v4
       with:
         repository: data-apis/array-api-tests
-        ref: '3cf8ef654c456d9fd1633d64e67b4470465940e9'  # Latest commit as of 2024-04-09
+        ref: '809a1984414cfc0bca68a823aeaeba7df3900d17'  # Latest commit as of 2024-06-26
         submodules: 'true'
         path: 'array-api-tests'
     - name: Set up Python

doc/release/upcoming_changes/26724.new_feature.rst

@@ -0,0 +1,7 @@
+* `numpy.cumulative_sum` and `numpy.cumulative_prod` were added as Array API
+  compatible alternatives for `numpy.cumsum` and `numpy.cumprod`. The new functions
+  can include a fixed initial (zeros for ``sum`` and ones for ``prod``) in the result.
+* `numpy.clip` now supports ``max`` and ``min`` keyword arguments which are meant
+  to replace ``a_min`` and ``a_max``. Also, for ``np.clip(a)`` or ``np.clip(a, None, None)``
+  a copy of the input array will be returned instead of raising an error.
+* `numpy.astype` now supports ``device`` argument.

doc/source/reference/routines.math.rst

@@ -63,6 +63,8 @@ Sums, products, differences
    sum
    nanprod
    nansum
+   cumulative_sum
+   cumulative_prod
    cumprod
    cumsum
    nancumprod

numpy/__init__.py

@@ -134,23 +134,23 @@
         can_cast, cbrt, cdouble, ceil, character, choose, clip, clongdouble,
         complex128, complex64, complexfloating, compress, concat, concatenate,
         conj, conjugate, convolve, copysign, copyto, correlate, cos, cosh,
-        count_nonzero, cross, csingle, cumprod, cumsum,
-        datetime64, datetime_as_string, datetime_data, deg2rad, degrees,
-        diagonal, divide, divmod, dot, double, dtype, e, einsum, einsum_path,
-        empty, empty_like, equal, errstate, euler_gamma, exp, exp2, expm1,
-        fabs, finfo, flatiter, flatnonzero, flexible, float16, float32,
-        float64, float_power, floating, floor, floor_divide, fmax, fmin, fmod,
-        format_float_positional, format_float_scientific, frexp, from_dlpack,
-        frombuffer, fromfile, fromfunction, fromiter, frompyfunc, fromstring,
-        full, full_like, gcd, generic, geomspace, get_printoptions,
-        getbufsize, geterr, geterrcall, greater, greater_equal, half,
-        heaviside, hstack, hypot, identity, iinfo, iinfo, indices, inexact,
-        inf, inner, int16, int32, int64, int8, int_, intc, integer, intp,
-        invert, is_busday, isclose, isdtype, isfinite, isfortran, isinf,
-        isnan, isnat, isscalar, issubdtype, lcm, ldexp, left_shift, less,
-        less_equal, lexsort, linspace, little_endian, log, log10, log1p, log2,
-        logaddexp, logaddexp2, logical_and, logical_not, logical_or,
-        logical_xor, logspace, long, longdouble, longlong, matmul,
+        count_nonzero, cross, csingle, cumprod, cumsum, cumulative_prod,
+        cumulative_sum, datetime64, datetime_as_string, datetime_data,
+        deg2rad, degrees, diagonal, divide, divmod, dot, double, dtype, e,
+        einsum, einsum_path, empty, empty_like, equal, errstate, euler_gamma,
+        exp, exp2, expm1, fabs, finfo, flatiter, flatnonzero, flexible,
+        float16, float32, float64, float_power, floating, floor, floor_divide,
+        fmax, fmin, fmod, format_float_positional, format_float_scientific,
+        frexp, from_dlpack, frombuffer, fromfile, fromfunction, fromiter,
+        frompyfunc, fromstring, full, full_like, gcd, generic, geomspace,
+        get_printoptions, getbufsize, geterr, geterrcall, greater,
+        greater_equal, half, heaviside, hstack, hypot, identity, iinfo, iinfo,
+        indices, inexact, inf, inner, int16, int32, int64, int8, int_, intc,
+        integer, intp, invert, is_busday, isclose, isdtype, isfinite,
+        isfortran, isinf, isnan, isnat, isscalar, issubdtype, lcm, ldexp,
+        left_shift, less, less_equal, lexsort, linspace, little_endian, log,
+        log10, log1p, log2, logaddexp, logaddexp2, logical_and, logical_not,
+        logical_or, logical_xor, logspace, long, longdouble, longlong, matmul,
         matrix_transpose, max, maximum, may_share_memory, mean, memmap, min,
         min_scalar_type, minimum, mod, modf, moveaxis, multiply, nan, ndarray,
         ndim, nditer, negative, nested_iters, newaxis, nextafter, nonzero,
@@ -290,7 +290,7 @@
     # import with `from numpy import *`.
     __future_scalars__ = {"str", "bytes", "object"}
 
-    __array_api_version__ = "2022.12"
+    __array_api_version__ = "2023.12"
 
     from ._array_api_info import __array_namespace_info__
 

numpy/__init__.pyi

@@ -260,13 +260,15 @@ from numpy._core.fromnumeric import (
     all as all,
     any as any,
     cumsum as cumsum,
+    cumulative_sum as cumulative_sum,
     ptp as ptp,
     max as max,
     min as min,
     amax as amax,
     amin as amin,
     prod as prod,
     cumprod as cumprod,
+    cumulative_prod as cumulative_prod,
     ndim as ndim,
     size as size,
     around as around,

numpy/_core/_methods.py

@@ -98,8 +98,8 @@ def _count_reduce_items(arr, axis, keepdims=False, where=True):
 
 def _clip(a, min=None, max=None, out=None, **kwargs):
     if min is None and max is None:
-        raise ValueError("One of max or min must be given")
-
+        # return identity
+        return um.positive(a, out=out, **kwargs)
     if min is None:
         return um.minimum(a, max, out=out, **kwargs)
     elif max is None:

numpy/_core/fromnumeric.py

@@ -21,8 +21,8 @@
 __all__ = [
     'all', 'amax', 'amin', 'any', 'argmax',
     'argmin', 'argpartition', 'argsort', 'around', 'choose', 'clip',
-    'compress', 'cumprod', 'cumsum', 'diagonal', 'mean',
-    'max', 'min', 'matrix_transpose',
+    'compress', 'cumprod', 'cumsum', 'cumulative_prod', 'cumulative_sum',
+    'diagonal', 'mean', 'max', 'min', 'matrix_transpose',
     'ndim', 'nonzero', 'partition', 'prod', 'ptp', 'put',
     'ravel', 'repeat', 'reshape', 'resize', 'round',
     'searchsorted', 'shape', 'size', 'sort', 'squeeze',
@@ -2210,12 +2210,14 @@ def compress(condition, a, axis=None, out=None):
     return _wrapfunc(a, 'compress', condition, axis=axis, out=out)
 
 
-def _clip_dispatcher(a, a_min, a_max, out=None, **kwargs):
-    return (a, a_min, a_max)
+def _clip_dispatcher(a, a_min=None, a_max=None, out=None, *, min=None,
+                     max=None, **kwargs):
+    return (a, a_min, a_max, out, min, max)
 
 
 @array_function_dispatch(_clip_dispatcher)
-def clip(a, a_min, a_max, out=None, **kwargs):
+def clip(a, a_min=np._NoValue, a_max=np._NoValue, out=None, *,
+         min=np._NoValue, max=np._NoValue, **kwargs):
     """
     Clip (limit) the values in an array.
 
@@ -2234,12 +2236,19 @@ def clip(a, a_min, a_max, out=None, **kwargs):
         Array containing elements to clip.
     a_min, a_max : array_like or None
         Minimum and maximum value. If ``None``, clipping is not performed on
-        the corresponding edge. Only one of `a_min` and `a_max` may be
-        ``None``. Both are broadcast against `a`.
+        the corresponding edge. If both ``a_min`` and ``a_max`` are ``None``,
+        the elements of the returned array stay the same. Both are broadcasted
+        against ``a``.
     out : ndarray, optional
         The results will be placed in this array. It may be the input
         array for in-place clipping.  `out` must be of the right shape
         to hold the output.  Its type is preserved.
+    min, max : array_like or None
+        Array API compatible alternatives for ``a_min`` and ``a_max``
+        arguments. Either ``a_min`` and ``a_max`` or ``min`` and ``max``
+        can be passed at the same time. Default: ``None``.
+
+        .. versionadded:: 2.1.0
     **kwargs
         For other keyword-only arguments, see the
         :ref:`ufunc docs <ufuncs.kwargs>`.
@@ -2283,6 +2292,19 @@ def clip(a, a_min, a_max, out=None, **kwargs):
     array([3, 4, 2, 3, 4, 5, 6, 7, 8, 8])
 
     """
+    if a_min is np._NoValue and a_max is np._NoValue:
+        a_min = None if min is np._NoValue else min
+        a_max = None if max is np._NoValue else max
+    elif a_min is np._NoValue:
+        raise TypeError("clip() missing 1 required positional "
+                        "argument: 'a_min'")
+    elif a_max is np._NoValue:
+        raise TypeError("clip() missing 1 required positional "
+                        "argument: 'a_max'")
+    elif min is not np._NoValue or max is not np._NoValue:
+        raise ValueError("Passing `min` or `max` keyword argument when "
+                         "`a_min` and `a_max` are provided is forbidden.")
+
     return _wrapfunc(a, 'clip', a_min, a_max, out=out, **kwargs)
 
 
@@ -2643,6 +2665,202 @@ def all(a, axis=None, out=None, keepdims=np._NoValue, *, where=np._NoValue):
                                   keepdims=keepdims, where=where)
 
 
+def _cumulative_func(x, func, axis, dtype, out, include_initial):
+    x = np.atleast_1d(x)
+    x_ndim = x.ndim
+    if axis is None:
+        if x_ndim >= 2:
+            raise ValueError("For arrays which have more than one dimension "
+                            "``axis`` argument is required.")
+        axis = 0
+
+    if out is not None and include_initial:
+        item = [slice(None)] * x_ndim
+        item[axis] = slice(1, None)
+        func.accumulate(x, axis=axis, dtype=dtype, out=out[tuple(item)])
+        item[axis] = 0
+        out[tuple(item)] = func.identity
+        return out
+
+    res = func.accumulate(x, axis=axis, dtype=dtype, out=out)
+    if include_initial:
+        initial_shape = list(x.shape)
+        initial_shape[axis] = 1
+        res = np.concat(
+            [np.full_like(res, func.identity, shape=initial_shape), res],
+            axis=axis,
+        )
+
+    return res
+
+
+def _cumulative_prod_dispatcher(x, /, *, axis=None, dtype=None, out=None,
+                                include_initial=None):
+    return (x, out)
+
+
+@array_function_dispatch(_cumulative_prod_dispatcher)
+def cumulative_prod(x, /, *, axis=None, dtype=None, out=None,
+                    include_initial=False):
+    """
+    Return the cumulative product of elements along a given axis.
+
+    This function is an Array API compatible alternative to `numpy.cumprod`.
+
+    Parameters
+    ----------
+    x : array_like
+        Input array.
+    axis : int, optional
+        Axis along which the cumulative product is computed. The default
+        (None) is only allowed for one-dimensional arrays. For arrays
+        with more than one dimension ``axis`` is required.
+    dtype : dtype, optional
+        Type of the returned array, as well as of the accumulator in which
+        the elements are multiplied.  If ``dtype`` is not specified, it
+        defaults to the dtype of ``x``, unless ``x`` has an integer dtype
+        with a precision less than that of the default platform integer.
+        In that case, the default platform integer is used instead.
+    out : ndarray, optional
+        Alternative output array in which to place the result. It must
+        have the same shape and buffer length as the expected output
+        but the type of the resulting values will be cast if necessary.
+        See :ref:`ufuncs-output-type` for more details.
+    include_initial : bool, optional
+        Boolean indicating whether to include the initial value (ones) as
+        the first value in the output. With ``include_initial=True``
+        the shape of the output is different than the shape of the input.
+        Default: ``False``.
+
+    Returns
+    -------
+    cumulative_prod_along_axis : ndarray
+        A new array holding the result is returned unless ``out`` is
+        specified, in which case a reference to ``out`` is returned. The
+        result has the same shape as ``x`` if ``include_initial=False``.
+
+    Notes
+    -----
+    Arithmetic is modular when using integer types, and no error is
+    raised on overflow.
+
+    Examples
+    --------
+    >>> a = np.array([1, 2, 3])
+    >>> np.cumulative_prod(a)  # intermediate results 1, 1*2
+    ...                        # total product 1*2*3 = 6
+    array([1, 2, 6])
+    >>> a = np.array([1, 2, 3, 4, 5, 6])
+    >>> np.cumulative_prod(a, dtype=float) # specify type of output
+    array([   1.,    2.,    6.,   24.,  120.,  720.])
+
+    The cumulative product for each column (i.e., over the rows) of ``b``:
+
+    >>> b = np.array([[1, 2, 3], [4, 5, 6]])
+    >>> np.cumulative_prod(b, axis=0)
+    array([[ 1,  2,  3],
+           [ 4, 10, 18]])
+
+    The cumulative product for each row (i.e. over the columns) of ``b``:
+
+    >>> np.cumulative_prod(b, axis=1)
+    array([[  1,   2,   6],
+           [  4,  20, 120]])
+
+    """
+    return _cumulative_func(x, um.multiply, axis, dtype, out, include_initial)
+
+
+def _cumulative_sum_dispatcher(x, /, *, axis=None, dtype=None, out=None,
+                               include_initial=None):
+    return (x, out)
+
+
+@array_function_dispatch(_cumulative_sum_dispatcher)
+def cumulative_sum(x, /, *, axis=None, dtype=None, out=None,
+                   include_initial=False):
+    """
+    Return the cumulative sum of the elements along a given axis.
+
+    This function is an Array API compatible alternative to `numpy.cumsum`.
+
+    Parameters
+    ----------
+    x : array_like
+        Input array.
+    axis : int, optional
+        Axis along which the cumulative sum is computed. The default
+        (None) is only allowed for one-dimensional arrays. For arrays
+        with more than one dimension ``axis`` is required.
+    dtype : dtype, optional
+        Type of the returned array and of the accumulator in which the
+        elements are summed.  If ``dtype`` is not specified, it defaults
+        to the dtype of ``x``, unless ``x`` has an integer dtype with
+        a precision less than that of the default platform integer.
+        In that case, the default platform integer is used.
+    out : ndarray, optional
+        Alternative output array in which to place the result. It must
+        have the same shape and buffer length as the expected output
+        but the type will be cast if necessary. See :ref:`ufuncs-output-type`
+        for more details.
+    include_initial : bool, optional
+        Boolean indicating whether to include the initial value (ones) as
+        the first value in the output. With ``include_initial=True``
+        the shape of the output is different than the shape of the input.
+        Default: ``False``.
+
+    Returns
+    -------
+    cumulative_sum_along_axis : ndarray
+        A new array holding the result is returned unless ``out`` is
+        specified, in which case a reference to ``out`` is returned. The
+        result has the same shape as ``x`` if ``include_initial=False``.
+
+    See Also
+    --------
+    sum : Sum array elements.
+    trapezoid : Integration of array values using composite trapezoidal rule.
+    diff : Calculate the n-th discrete difference along given axis.
+
+    Notes
+    -----
+    Arithmetic is modular when using integer types, and no error is
+    raised on overflow.
+
+    ``cumulative_sum(a)[-1]`` may not be equal to ``sum(a)`` for
+    floating-point values since ``sum`` may use a pairwise summation routine,
+    reducing the roundoff-error. See `sum` for more information.
+
+    Examples
+    --------
+    >>> a = np.array([1, 2, 3, 4, 5, 6])
+    >>> a
+    array([1, 2, 3, 4, 5, 6])
+    >>> np.cumulative_sum(a)
+    array([ 1,  3,  6, 10, 15, 21])
+    >>> np.cumulative_sum(a, dtype=float)  # specifies type of output value(s)
+    array([  1.,   3.,   6.,  10.,  15.,  21.])
+
+    >>> b = np.array([[1, 2, 3], [4, 5, 6]])
+    >>> np.cumulative_sum(b,axis=0)  # sum over rows for each of the 3 columns
+    array([[1, 2, 3],
+           [5, 7, 9]])
+    >>> np.cumulative_sum(b,axis=1)  # sum over columns for each of the 2 rows
+    array([[ 1,  3,  6],
+           [ 4,  9, 15]])
+
+    ``cumulative_sum(c)[-1]`` may not be equal to ``sum(c)``
+
+    >>> c = np.array([1, 2e-9, 3e-9] * 1000000)
+    >>> np.cumulative_sum(c)[-1]
+    1000000.0050045159
+    >>> c.sum()
+    1000000.0050000029
+
+    """
+    return _cumulative_func(x, um.add, axis, dtype, out, include_initial)
+
+
 def _cumsum_dispatcher(a, axis=None, dtype=None, out=None):
     return (a, out)
 
@@ -2681,6 +2899,7 @@ def cumsum(a, axis=None, dtype=None, out=None):
 
     See Also
     --------
+    cumulative_sum : Array API compatible alternative for ``cumsum``.
     sum : Sum array elements.
     trapezoid : Integration of array values using composite trapezoidal rule.
     diff : Calculate the n-th discrete difference along given axis.
@@ -3269,6 +3488,7 @@ def cumprod(a, axis=None, dtype=None, out=None):
 
     See Also
     --------
+    cumulative_prod : Array API compatible alternative for ``cumprod``.
     :ref:`ufuncs-output-type`
 
     Notes