Implement dpnp.nancumsum() through dpnp.cumsum call (#1781)

* Implement dpnp.nancumsum() through existing calls

* Enabled strict dtype check due to recent changes in dpnp.sum()

* FFT tests still require check_only_type_kind=True

Author: antonwolfy

doc/reference/math.rst

@@ -66,17 +66,17 @@ Sums, products, differences
 
    dpnp.prod
    dpnp.sum
+   dpnp.nanprod
+   dpnp.nansum
    dpnp.cumprod
    dpnp.cumsum
    dpnp.nancumprod
    dpnp.nancumsum
-   dpnp.nansum
-   dpnp.nanprod
-   dpnp.cross
    dpnp.diff
    dpnp.ediff1d
    dpnp.gradient
    dpnp.trapz
+   dpnp.cross
 
 
 Exponents and logarithms

dpnp/backend/include/dpnp_iface_fptr.hpp

@@ -108,8 +108,6 @@ enum class DPNPFuncName : size_t
     DPNP_FN_CUMPROD_EXT,   /**< Used in numpy.cumprod() impl, requires extra
                               parameters */
     DPNP_FN_CUMSUM,        /**< Used in numpy.cumsum() impl  */
-    DPNP_FN_CUMSUM_EXT,    /**< Used in numpy.cumsum() impl, requires extra
-                              parameters */
     DPNP_FN_DEGREES,       /**< Used in numpy.degrees() impl  */
     DPNP_FN_DEGREES_EXT,   /**< Used in numpy.degrees() impl, requires extra
                               parameters */

dpnp/backend/kernels/dpnp_krnl_mathematical.cpp

@@ -425,14 +425,6 @@ template <typename _DataType_input, typename _DataType_output>
 void (*dpnp_cumsum_default_c)(void *, void *, size_t) =
     dpnp_cumsum_c<_DataType_input, _DataType_output>;
 
-template <typename _DataType_input, typename _DataType_output>
-DPCTLSyclEventRef (*dpnp_cumsum_ext_c)(DPCTLSyclQueueRef,
-                                       void *,
-                                       void *,
-                                       size_t,
-                                       const DPCTLEventVectorRef) =
-    dpnp_cumsum_c<_DataType_input, _DataType_output>;
-
 template <typename _KernelNameSpecialization1,
           typename _KernelNameSpecialization2>
 class dpnp_ediff1d_c_kernel;
@@ -1179,15 +1171,6 @@ void func_map_init_mathematical(func_map_t &fmap)
     fmap[DPNPFuncName::DPNP_FN_CUMSUM][eft_DBL][eft_DBL] = {
         eft_DBL, (void *)dpnp_cumsum_default_c<double, double>};
 
-    fmap[DPNPFuncName::DPNP_FN_CUMSUM_EXT][eft_INT][eft_INT] = {
-        eft_LNG, (void *)dpnp_cumsum_ext_c<int32_t, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_CUMSUM_EXT][eft_LNG][eft_LNG] = {
-        eft_LNG, (void *)dpnp_cumsum_ext_c<int64_t, int64_t>};
-    fmap[DPNPFuncName::DPNP_FN_CUMSUM_EXT][eft_FLT][eft_FLT] = {
-        eft_FLT, (void *)dpnp_cumsum_ext_c<float, float>};
-    fmap[DPNPFuncName::DPNP_FN_CUMSUM_EXT][eft_DBL][eft_DBL] = {
-        eft_DBL, (void *)dpnp_cumsum_ext_c<double, double>};
-
     fmap[DPNPFuncName::DPNP_FN_EDIFF1D][eft_INT][eft_INT] = {
         eft_LNG, (void *)dpnp_ediff1d_default_c<int32_t, int64_t>};
     fmap[DPNPFuncName::DPNP_FN_EDIFF1D][eft_LNG][eft_LNG] = {

dpnp/dpnp_algo/dpnp_algo.pxd

@@ -38,7 +38,6 @@ cdef extern from "dpnp_iface_fptr.hpp" namespace "DPNPFuncName":  # need this na
         DPNP_FN_COPY_EXT
         DPNP_FN_CORRELATE_EXT
         DPNP_FN_CUMPROD_EXT
-        DPNP_FN_CUMSUM_EXT
         DPNP_FN_DEGREES_EXT
         DPNP_FN_DIAG_INDICES_EXT
         DPNP_FN_DIAGONAL_EXT

dpnp/dpnp_algo/dpnp_algo_mathematical.pxi

@@ -45,7 +45,6 @@ __all__ += [
     "dpnp_fmin",
     "dpnp_modf",
     "dpnp_nancumprod",
-    "dpnp_nancumsum",
     "dpnp_trapz",
 ]
 
@@ -70,18 +69,6 @@ cpdef utils.dpnp_descriptor dpnp_cumprod(utils.dpnp_descriptor x1):
     return call_fptr_1in_1out(DPNP_FN_CUMPROD_EXT, x1, (x1.size,))
 
 
-cpdef utils.dpnp_descriptor dpnp_cumsum(utils.dpnp_descriptor x1):
-    # instead of x1.shape, (x1.size, ) is passed to the function
-    # due to the following:
-    # >>> import numpy
-    # >>> a = numpy.array([[1, 2], [2, 3]])
-    # >>> res = numpy.cumsum(a)
-    # >>> res.shape
-    # (4,)
-
-    return call_fptr_1in_1out(DPNP_FN_CUMSUM_EXT, x1, (x1.size,))
-
-
 cpdef utils.dpnp_descriptor dpnp_ediff1d(utils.dpnp_descriptor x1):
 
     if x1.size <= 1:
@@ -253,19 +240,6 @@ cpdef utils.dpnp_descriptor dpnp_nancumprod(utils.dpnp_descriptor x1):
     return dpnp_cumprod(x1_desc)
 
 
-cpdef utils.dpnp_descriptor dpnp_nancumsum(utils.dpnp_descriptor x1):
-    cur_x1 = x1.get_pyobj().copy()
-
-    cur_x1_flatiter = cur_x1.flat
-
-    for i in range(cur_x1.size):
-        if dpnp.isnan(cur_x1_flatiter[i]):
-            cur_x1_flatiter[i] = 0
-
-    x1_desc = dpnp.get_dpnp_descriptor(cur_x1, copy_when_nondefault_queue=False)
-    return dpnp_cumsum(x1_desc)
-
-
 cpdef utils.dpnp_descriptor dpnp_trapz(utils.dpnp_descriptor y1, utils.dpnp_descriptor x1, double dx):
 
     cdef DPNPFuncType param1_type = dpnp_dtype_to_DPNPFuncType(y1.dtype)

dpnp/dpnp_iface_nanfunctions.py

@@ -46,7 +46,6 @@
 # pylint: disable=no-name-in-module
 from .dpnp_algo import (
     dpnp_nancumprod,
-    dpnp_nancumsum,
 )
 from .dpnp_utils import (
     call_origin,
@@ -292,43 +291,68 @@ def nancumprod(x1, **kwargs):
     return call_origin(numpy.nancumprod, x1, **kwargs)
 
 
-def nancumsum(x1, **kwargs):
+def nancumsum(a, axis=None, dtype=None, out=None):
     """
-    Return the cumulative sum of the elements along a given axis.
+    Return the cumulative sum of array elements over a given axis treating
+    Not a Numbers (NaNs) as zero. The cumulative sum does not change when NaNs
+    are encountered and leading NaNs are replaced by zeros.
 
     For full documentation refer to :obj:`numpy.nancumsum`.
 
-    Limitations
-    -----------
-    Parameter `x` is supported as :class:`dpnp.ndarray`.
-    Keyword argument `kwargs` is currently unsupported.
-    Otherwise the function will be executed sequentially on CPU.
-    Input array data types are limited by supported DPNP :ref:`Data types`.
+    Parameters
+    ----------
+    a : {dpnp.ndarray, usm_ndarray}
+        Input array.
+    axis : int, optional
+        Axis along which the cumulative sum is computed. The default (``None``)
+        is to compute the cumsum over the flattened array.
+    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
+        `a`, unless `a` 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 : {dpnp.ndarray, usm_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.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        A new array holding the result is returned unless `out` is specified as
+        :class:`dpnp.ndarray`, in which case a reference to `out` is returned.
+        The result has the same size as `a`, and the same shape as `a` if `axis`
+        is not ``None`` or `a` is a 1-d array.
 
     See Also
     --------
-    :obj:`dpnp.cumsum` : Return the cumulative sum of the elements
-                         along a given axis.
+    :obj:`dpnp.cumsum` : Cumulative sum across array propagating NaNs.
+    :obj:`dpnp.isnan` : Show which elements are NaN.
 
     Examples
     --------
     >>> import dpnp as np
-    >>> a = np.array([1., np.nan])
-    >>> result = np.nancumsum(a)
-    >>> [x for x in result]
-    [1.0, 1.0]
-    >>> b = np.array([[1., 2., np.nan], [4., np.nan, 6.]])
-    >>> result = np.nancumprod(b)
-    >>> [x for x in result]
-    [1.0, 3.0, 3.0, 7.0, 7.0, 13.0]
+    >>> np.nancumsum(np.array(1))
+    array([1])
+    >>> np.nancumsum(np.array([1]))
+    array([1])
+    >>> np.nancumsum(np.array([1, np.nan]))
+    array([1., 1.])
+    >>> a = np.array([[1, 2], [3, np.nan]])
+    >>> np.nancumsum(a)
+    array([1., 3., 6., 6.])
+    >>> np.nancumsum(a, axis=0)
+    array([[1., 2.],
+           [4., 2.]])
+    >>> np.nancumsum(a, axis=1)
+    array([[1., 3.],
+           [3., 3.]])
 
     """
 
-    x1_desc = dpnp.get_dpnp_descriptor(x1, copy_when_nondefault_queue=False)
-    if x1_desc and not kwargs:
-        return dpnp_nancumsum(x1_desc).get_pyobj()
-
-    return call_origin(numpy.nancumsum, x1, **kwargs)
+    a, _ = _replace_nan(a, 0)
+    return dpnp.cumsum(a, axis=axis, dtype=dtype, out=out)
 
 
 def nanmax(a, axis=None, out=None, keepdims=False, initial=None, where=True):

tests/test_linalg.py

@@ -961,13 +961,10 @@ def test_norm_0D(self, ord, axis):
     def test_norm_1D(self, dtype, ord, axis, keepdims):
         a = numpy.array(numpy.random.uniform(-5, 5, 10), dtype=dtype)
         ia = inp.array(a)
+
         result = inp.linalg.norm(ia, ord=ord, axis=axis, keepdims=keepdims)
         expected = numpy.linalg.norm(a, ord=ord, axis=axis, keepdims=keepdims)
-        # use only type kinds check when dpnp handles complex64 arrays
-        # since `dpnp.sum()` and `numpy.sum()` return different dtypes
-        assert_dtype_allclose(
-            result, expected, check_only_type_kind=(dtype == inp.float32)
-        )
+        assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_complex_dtypes())
@@ -983,11 +980,10 @@ def test_norm_1D_complex(self, dtype, ord, axis, keepdims):
         x2 = numpy.random.uniform(-5, 5, 10)
         a = numpy.array(x1 + 1j * x2, dtype=dtype)
         ia = inp.array(a)
+
         result = inp.linalg.norm(ia, ord=ord, axis=axis, keepdims=keepdims)
         expected = numpy.linalg.norm(a, ord=ord, axis=axis, keepdims=keepdims)
-        assert_dtype_allclose(
-            result, expected, check_only_type_kind=(dtype == inp.complex64)
-        )
+        assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True))
@@ -1027,9 +1023,7 @@ def test_norm_2D(self, dtype, ord, axis, keepdims):
             expected = numpy.linalg.norm(
                 a, ord=ord, axis=axis, keepdims=keepdims
             )
-            assert_dtype_allclose(
-                result, expected, check_only_type_kind=(dtype == inp.float32)
-            )
+            assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_complex_dtypes())
@@ -1069,9 +1063,7 @@ def test_norm_2D_complex(self, dtype, ord, axis, keepdims):
             expected = numpy.linalg.norm(
                 a, ord=ord, axis=axis, keepdims=keepdims
             )
-            assert_dtype_allclose(
-                result, expected, check_only_type_kind=(dtype == inp.complex64)
-            )
+            assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True))
@@ -1117,9 +1109,7 @@ def test_norm_ND(self, dtype, ord, axis, keepdims):
             expected = numpy.linalg.norm(
                 a, ord=ord, axis=axis, keepdims=keepdims
             )
-            assert_dtype_allclose(
-                result, expected, check_only_type_kind=(dtype == inp.float32)
-            )
+            assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_complex_dtypes())
@@ -1165,9 +1155,7 @@ def test_norm_ND_complex(self, dtype, ord, axis, keepdims):
             expected = numpy.linalg.norm(
                 a, ord=ord, axis=axis, keepdims=keepdims
             )
-            assert_dtype_allclose(
-                result, expected, check_only_type_kind=(dtype == inp.complex64)
-            )
+            assert_dtype_allclose(result, expected)
 
     @pytest.mark.usefixtures("suppress_divide_numpy_warnings")
     @pytest.mark.parametrize("dtype", get_all_dtypes())
@@ -1213,11 +1201,7 @@ def test_norm_usm_ndarray(self, dtype, ord, axis, keepdims):
             expected = numpy.linalg.norm(
                 a, ord=ord, axis=axis, keepdims=keepdims
             )
-            assert_dtype_allclose(
-                result,
-                expected,
-                check_only_type_kind=(dtype in [inp.float32, inp.complex64]),
-            )
+            assert_dtype_allclose(result, expected)
 
     @pytest.mark.parametrize("stride", [3, -1, -5], ids=["3", "-1", "-5"])
     def test_norm_strided_1D(self, stride):