diff --git a/doc/reference/math.rst b/doc/reference/math.rst index bc4747429f9b..48142b10fd45 100644 --- a/doc/reference/math.rst +++ b/doc/reference/math.rst @@ -96,6 +96,7 @@ Exponents and logarithms dpnp.logaddexp dpnp.logaddexp2 dpnp.logsumexp + dpnp.cumlogsumexp Other special functions diff --git a/dpnp/dpnp_iface_trigonometric.py b/dpnp/dpnp_iface_trigonometric.py index 5aa8bcb798c5..64c110190bf8 100644 --- a/dpnp/dpnp_iface_trigonometric.py +++ b/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. diff --git a/tests/test_mathematical.py b/tests/test_mathematical.py index e9e9214cd4bd..84474f27823e 100644 --- a/tests/test_mathematical.py +++ b/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", diff --git a/tests/test_strides.py b/tests/test_strides.py index 6fe28281fc62..780c069e9ca0 100644 --- a/tests/test_strides.py +++ b/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] diff --git a/tests/test_sycl_queue.py b/tests/test_sycl_queue.py index 8080868cff4d..88ca87d833ea 100644 --- a/tests/test_sycl_queue.py +++ b/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, diff --git a/tests/test_usm_type.py b/tests/test_usm_type.py index 278327ff6a45..411b91684911 100644 --- a/tests/test_usm_type.py +++ b/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]]]),