Leverage dpnp.cumlogsumexp through dpctl.tensor implementation (#1816)

* Implement dpnp.cumprod through dpctl.tensor

* Implement dpnp.nancumprod() through existing calls

* Implement dpnp.cumlogsumexp through dpctl.tensor

* Resolved pre-commit issues

* Applied review comment

* Fix test_out running on CPU

* Generate docstings documentation for dpnp.cumlogsumexp

Author: antonwolfy

doc/reference/math.rst

@@ -96,6 +96,7 @@ Exponents and logarithms
    dpnp.logaddexp
    dpnp.logaddexp2
    dpnp.logsumexp
+   dpnp.cumlogsumexp
 
 
 Other special functions

dpnp/dpnp_iface_trigonometric.py

@@ -71,6 +71,7 @@
     "cbrt",
     "cos",
     "cosh",
+    "cumlogsumexp",
     "deg2rad",
     "degrees",
     "exp",
@@ -665,6 +666,94 @@ def _get_accumulation_res_dt(a, dtype, _out):
 )
 
 
+def cumlogsumexp(
+    x, /, *, axis=None, dtype=None, include_initial=False, out=None
+):
+    """
+    Calculates the cumulative logarithm of the sum of elements in the input
+    array `x`.
+
+    Parameters
+    ----------
+    x : {dpnp.ndarray, usm_ndarray}
+        Input array, expected to have a real-valued data type.
+    axis : {None, int}, optional
+        Axis or axes along which values must be computed. If a tuple of unique
+        integers, values are computed over multiple axes. If ``None``, the
+        result is computed over the entire array.
+        Default: ``None``.
+    dtype : {None, dtype}, optional
+        Data type of the returned array. If ``None``, the default data type is
+        inferred from the "kind" of the input array data type.
+
+        - If `x` has a real-valued floating-point data type, the returned array
+          will have the same data type as `x`.
+        - If `x` has a boolean or integral data type, the returned array will
+          have the default floating point data type for the device where input
+          array `x` is allocated.
+        - If `x` has a complex-valued floating-point data type, an error is
+          raised.
+
+        If the data type (either specified or resolved) differs from the data
+        type of `x`, the input array elements are cast to the specified data
+        type before computing the result.
+        Default: ``None``.
+    include_initial : {None, bool}, optional
+        A boolean indicating whether to include the initial value (i.e., the
+        additive identity, zero) as the first value along the provided axis in
+        the output.
+        Default: ``False``.
+    out : {None, dpnp.ndarray, usm_ndarray}, optional
+        The array into which the result is written. The data type of `out` must
+        match the expected shape and the expected data type of the result or
+        (if provided) `dtype`. If ``None`` then a new array is returned.
+        Default: ``None``.
+
+    Returns
+    -------
+    out : dpnp.ndarray
+        An array containing the results. If the result was computed over the
+        entire array, a zero-dimensional array is returned. The returned array
+        has the data type as described in the `dtype` parameter description
+        above.
+
+    Note
+    ----
+    This function is equivalent of `numpy.logaddexp.accumulate`.
+
+    See Also
+    --------
+    :obj:`dpnp.logsumexp` : Logarithm of the sum of elements of the inputs,
+                            element-wise.
+
+    Examples
+    --------
+    >>> import dpnp as np
+    >>> a = np.ones(10)
+    >>> np.cumlogsumexp(a)
+    array([1.        , 1.69314718, 2.09861229, 2.38629436, 2.60943791,
+           2.79175947, 2.94591015, 3.07944154, 3.19722458, 3.30258509])
+
+    """
+
+    dpnp.check_supported_arrays_type(x)
+    if x.ndim > 1 and axis is None:
+        usm_x = dpnp.ravel(x).get_array()
+    else:
+        usm_x = dpnp.get_usm_ndarray(x)
+
+    return dpnp_wrap_reduction_call(
+        x,
+        out,
+        dpt.cumulative_logsumexp,
+        _get_accumulation_res_dt,
+        usm_x,
+        axis=axis,
+        dtype=dtype,
+        include_initial=include_initial,
+    )
+
+
 def deg2rad(x1):
     """
     Convert angles from degrees to radians.

tests/test_mathematical.py

@@ -161,6 +161,114 @@ def test_not_implemented_kwargs(self, kwargs):
             dpnp.clip(a, 1, 5, **kwargs)
 
 
+class TestCumLogSumExp:
+    def _assert_arrays(self, res, exp, axis, include_initial):
+        if include_initial:
+            if axis != None:
+                res_initial = dpnp.take(res, dpnp.array([0]), axis=axis)
+                res_no_initial = dpnp.take(
+                    res, dpnp.array(range(1, res.shape[axis])), axis=axis
+                )
+            else:
+                res_initial = res[0]
+                res_no_initial = res[1:]
+            assert_dtype_allclose(res_no_initial, exp)
+            assert (res_initial == -dpnp.inf).all()
+        else:
+            assert_dtype_allclose(res, exp)
+
+    def _get_exp_array(self, a, axis, dtype):
+        np_a = dpnp.asnumpy(a)
+        if axis != None:
+            return numpy.logaddexp.accumulate(np_a, axis=axis, dtype=dtype)
+        return numpy.logaddexp.accumulate(np_a.ravel(), dtype=dtype)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True))
+    @pytest.mark.parametrize("axis", [None, 2, -1])
+    @pytest.mark.parametrize("include_initial", [True, False])
+    def test_basic(self, dtype, axis, include_initial):
+        a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
+        res = dpnp.cumlogsumexp(a, axis=axis, include_initial=include_initial)
+
+        exp_dt = None
+        if dtype == dpnp.bool:
+            exp_dt = dpnp.default_float_type(a.device)
+
+        exp = self._get_exp_array(a, axis, exp_dt)
+        self._assert_arrays(res, exp, axis, include_initial)
+
+    @pytest.mark.parametrize("dtype", get_all_dtypes(no_complex=True))
+    @pytest.mark.parametrize("axis", [None, 2, -1])
+    @pytest.mark.parametrize("include_initial", [True, False])
+    def test_out(self, dtype, axis, include_initial):
+        a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
+
+        if dpnp.issubdtype(a, dpnp.float32):
+            exp_dt = dpnp.float32
+        else:
+            exp_dt = dpnp.default_float_type(a.device)
+
+        if axis != None:
+            if include_initial:
+                norm_axis = numpy.core.numeric.normalize_axis_index(
+                    axis, a.ndim, "axis"
+                )
+                out_sh = (
+                    a.shape[:norm_axis]
+                    + (a.shape[norm_axis] + 1,)
+                    + a.shape[norm_axis + 1 :]
+                )
+            else:
+                out_sh = a.shape
+        else:
+            out_sh = (a.size + int(include_initial),)
+        out = dpnp.empty_like(a, shape=out_sh, dtype=exp_dt)
+        res = dpnp.cumlogsumexp(
+            a, axis=axis, include_initial=include_initial, out=out
+        )
+
+        exp = self._get_exp_array(a, axis, exp_dt)
+
+        assert res is out
+        self._assert_arrays(res, exp, axis, include_initial)
+
+    def test_axis_tuple(self):
+        a = dpnp.ones((3, 4))
+        assert_raises(TypeError, dpnp.cumlogsumexp, a, axis=(0, 1))
+
+    @pytest.mark.parametrize(
+        "in_dtype", get_all_dtypes(no_bool=True, no_complex=True)
+    )
+    @pytest.mark.parametrize("out_dtype", get_all_dtypes(no_bool=True))
+    def test_dtype(self, in_dtype, out_dtype):
+        a = dpnp.ones(100, dtype=in_dtype)
+        res = dpnp.cumlogsumexp(a, dtype=out_dtype)
+        exp = numpy.logaddexp.accumulate(dpnp.asnumpy(a))
+        exp = exp.astype(out_dtype)
+
+        assert_allclose(res, exp, rtol=1e-06)
+
+    @pytest.mark.usefixtures("suppress_invalid_numpy_warnings")
+    @pytest.mark.parametrize(
+        "arr_dt", get_all_dtypes(no_none=True, no_complex=True)
+    )
+    @pytest.mark.parametrize(
+        "out_dt", get_all_dtypes(no_none=True, no_complex=True)
+    )
+    @pytest.mark.parametrize("dtype", get_all_dtypes())
+    def test_out_dtype(self, arr_dt, out_dt, dtype):
+        a = numpy.arange(10, 20).reshape((2, 5)).astype(dtype=arr_dt)
+        out = numpy.zeros_like(a, dtype=out_dt)
+
+        ia = dpnp.array(a)
+        iout = dpnp.array(out)
+
+        result = dpnp.cumlogsumexp(ia, out=iout, dtype=dtype, axis=1)
+        exp = numpy.logaddexp.accumulate(a, out=out, axis=1)
+        assert_allclose(result, exp.astype(dtype), rtol=1e-06)
+        assert result is iout
+
+
 class TestCumProd:
     @pytest.mark.parametrize(
         "arr, axis",

tests/test_strides.py

@@ -125,6 +125,16 @@ def test_logsumexp(dtype):
     assert_allclose(result, expected)
 
 
+@pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True, no_complex=True))
+def test_cumlogsumexp(dtype):
+    a = numpy.arange(10, dtype=dtype)[::2]
+    dpa = dpnp.arange(10, dtype=dtype)[::2]
+
+    result = dpnp.cumlogsumexp(dpa)
+    expected = numpy.logaddexp.accumulate(a)
+    assert_allclose(result, expected)
+
+
 @pytest.mark.parametrize("dtype", get_all_dtypes(no_bool=True, no_complex=True))
 def test_reduce_hypot(dtype):
     a = numpy.arange(10, dtype=dtype)[::2]

tests/test_sycl_queue.py

@@ -556,6 +556,22 @@ def test_logsumexp(device):
     assert_sycl_queue_equal(result_queue, expected_queue)
 
 
+@pytest.mark.parametrize(
+    "device",
+    valid_devices,
+    ids=[device.filter_string for device in valid_devices],
+)
+def test_cumlogsumexp(device):
+    x = dpnp.arange(10, device=device)
+    result = dpnp.cumlogsumexp(x)
+    expected = numpy.logaddexp.accumulate(x.asnumpy())
+    assert_dtype_allclose(result, expected)
+
+    expected_queue = x.get_array().sycl_queue
+    result_queue = result.get_array().sycl_queue
+    assert_sycl_queue_equal(result_queue, expected_queue)
+
+
 @pytest.mark.parametrize(
     "device",
     valid_devices,

tests/test_usm_type.py

@@ -522,6 +522,7 @@ def test_norm(usm_type, ord, axis):
         ),
         pytest.param("cosh", [-5.0, -3.5, 0.0, 3.5, 5.0]),
         pytest.param("count_nonzero", [0, 1, 7, 0]),
+        pytest.param("cumlogsumexp", [1.0, 2.0, 4.0, 7.0]),
         pytest.param("cumprod", [[1, 2, 3], [4, 5, 6]]),
         pytest.param("cumsum", [[1, 2, 3], [4, 5, 6]]),
         pytest.param("diagonal", [[[1, 2], [3, 4]]]),