Align with new array api version in dpctl (#1774)

* Use utility function from dpctl.tensor._type_utils * Align sum and prod tests * Align logsumexp and reduce_hypot tests * Updated docstring of impacted functions
IntelPython · Apr 5, 2024 · 5a3e078 · 5a3e078
1 parent 42f2644
commit 5a3e078
Show file tree

Hide file tree

Showing 7 changed files with 72 additions and 125 deletions.
diff --git a/dpnp/dpnp_iface_mathematical.py b/dpnp/dpnp_iface_mathematical.py
@@ -39,6 +39,7 @@
 
 
 import dpctl.tensor as dpt
+import dpctl.tensor._type_utils as dtu
 import numpy
 from numpy.core.numeric import (
     normalize_axis_index,
@@ -2799,25 +2800,10 @@ def sum(
         If ``None``, the sum is computed over the entire array.
         Default: ``None``.
     dtype : 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 `a` has a real-valued floating-point data type,
-              the returned array will have the default real-valued
-              floating-point data type for the device where input
-              array `a` is allocated.
-            * If `a` has signed integral data type, the returned array
-              will have the default signed integral type for the device
-              where input array `a` is allocated.
-            * If `a` has unsigned integral data type, the returned array
-              will have the default unsigned integral type for the device
-              where input array `a` is allocated.
-            * If `a` has a complex-valued floating-point data type,
-              the returned array will have the default complex-valued
-              floating-pointer data type for the device where input
-              array `a` is allocated.
-            * If `a` has a boolean data type, the returned array will
-              have the default signed integral type for the device
-              where input array `a` is allocated.
+        Data type of the returned array. If ``None``, 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.
         If the data type (either specified or resolved) differs from the
         data type of `a`, the input array elements are cast to the
         specified data type before computing the sum.
@@ -2905,8 +2891,6 @@ def sum(
                 )
             )
         ):
-            from dpctl.tensor._reduction import _default_reduction_dtype
-
             from dpnp.backend.extensions.sycl_ext import _sycl_ext_impl
 
             input = a
@@ -2916,7 +2900,7 @@ def sum(
 
             queue = input.sycl_queue
             out_dtype = (
-                _default_reduction_dtype(input.dtype, queue)
+                dtu._default_accumulation_dtype(input.dtype, queue)
                 if dtype is None
                 else dtype
             )

diff --git a/dpnp/dpnp_iface_nanfunctions.py b/dpnp/dpnp_iface_nanfunctions.py
@@ -714,25 +714,10 @@ def nansum(
         If ``None``, the sum is computed over the entire array.
         Default: ``None``.
     dtype : 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 `a` has a real-valued floating-point data type,
-              the returned array will have the default real-valued
-              floating-point data type for the device where input
-              array `a` is allocated.
-            * If `a` has signed integral data type, the returned array
-              will have the default signed integral type for the device
-              where input array `a` is allocated.
-            * If `a` has unsigned integral data type, the returned array
-              will have the default unsigned integral type for the device
-              where input array `a` is allocated.
-            * If `a` has a complex-valued floating-point data type,
-              the returned array will have the default complex-valued
-              floating-pointer data type for the device where input
-              array `a` is allocated.
-            * If `a` has a boolean data type, the returned array will
-              have the default signed integral type for the device
-              where input array `a` is allocated.
+        Data type of the returned array. If ``None``, 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.
         If the data type (either specified or resolved) differs from the
         data type of `a`, the input array elements are cast to the
         specified data type before computing the sum.

diff --git a/dpnp/dpnp_iface_trigonometric.py b/dpnp/dpnp_iface_trigonometric.py
@@ -1352,20 +1352,14 @@ def logsumexp(x, axis=None, out=None, dtype=None, keepdims=False):
         If provided, the result will be inserted into this array. It should
         be of the appropriate shape and dtype.
     dtype : data type, 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 default real-valued
-              floating-point data type for the device where input
-              array `x` is allocated.
-            * 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.
+        Data type of the returned array. If ``None``, 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.
         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``.
+        specified data type before computing the result.
+        Default: ``None``.
     keepdims : bool
         If ``True``, the reduced axes (dimensions) are included in the result
         as singleton dimensions, so that the returned array remains
@@ -1498,20 +1492,14 @@ def reduce_hypot(x, axis=None, out=None, dtype=None, keepdims=False):
         If provided, the result will be inserted into this array. It should
         be of the appropriate shape and dtype.
     dtype : data type, 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 default real-valued
-                  floating-point data type for the device where input
-                  array `x` is allocated.
-                * 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``.
+        Data type of the returned array. If ``None``, 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.
+        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``.
     keepdims : bool
         If ``True``, the reduced axes (dimensions) are included in the result
         as singleton dimensions, so that the returned array remains

diff --git a/tests/test_arithmetic.py b/tests/test_arithmetic.py
@@ -6,15 +6,6 @@
 from tests.third_party.cupy import testing
 
 
-# Note: numpy.sum() always upcast integers to (u)int64 and float32 to
-# float64 for dtype=None. `np.sum` does that too for integers, but not for
-# float32, so we need to special-case it for these tests
-def _get_dtype_kwargs(xp, dtype):
-    if xp is numpy and dtype == numpy.float32 and has_support_aspect64():
-        return {"dtype": numpy.float64}
-    return {}
-
-
 class TestArithmetic(unittest.TestCase):
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose()
@@ -42,7 +33,7 @@ def test_nanprod(self, xp, dtype):
     @testing.numpy_cupy_allclose()
     def test_nansum(self, xp, dtype):
         a = xp.array([-2.5, -1.5, xp.nan, 10.5, 1.5, xp.nan], dtype=dtype)
-        return xp.nansum(a, **_get_dtype_kwargs(xp, a.dtype))
+        return xp.nansum(a)
 
     @testing.for_float_dtypes()
     @testing.numpy_cupy_allclose()

diff --git a/tests/test_mathematical.py b/tests/test_mathematical.py
@@ -1884,7 +1884,9 @@ class TestLogSumExp:
     def test_logsumexp(self, dtype, axis, keepdims):
         a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
         res = dpnp.logsumexp(a, axis=axis, keepdims=keepdims)
-        exp_dtype = dpnp.default_float_type(a.device)
+        exp_dtype = (
+            dpnp.default_float_type(a.device) if dtype == dpnp.bool else None
+        )
         exp = numpy.logaddexp.reduce(
             dpnp.asnumpy(a), axis=axis, keepdims=keepdims, dtype=exp_dtype
         )
@@ -1896,11 +1898,17 @@ def test_logsumexp(self, dtype, axis, keepdims):
     @pytest.mark.parametrize("keepdims", [True, False])
     def test_logsumexp_out(self, dtype, axis, keepdims):
         a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
-        exp_dtype = dpnp.default_float_type(a.device)
+        exp_dtype = (
+            dpnp.default_float_type(a.device) if dtype == dpnp.bool else None
+        )
         exp = numpy.logaddexp.reduce(
             dpnp.asnumpy(a), axis=axis, keepdims=keepdims, dtype=exp_dtype
         )
-        dpnp_out = dpnp.empty(exp.shape, dtype=exp_dtype)
+
+        exp_dtype = exp.dtype
+        if exp_dtype == numpy.float64 and not has_support_aspect64():
+            exp_dtype = numpy.float32
+        dpnp_out = dpnp.empty_like(a, shape=exp.shape, dtype=exp_dtype)
         res = dpnp.logsumexp(a, axis=axis, out=dpnp_out, keepdims=keepdims)
 
         assert res is dpnp_out
@@ -1926,7 +1934,9 @@ class TestReduceHypot:
     def test_reduce_hypot(self, dtype, axis, keepdims):
         a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
         res = dpnp.reduce_hypot(a, axis=axis, keepdims=keepdims)
-        exp_dtype = dpnp.default_float_type(a.device)
+        exp_dtype = (
+            dpnp.default_float_type(a.device) if dtype == dpnp.bool else None
+        )
         exp = numpy.hypot.reduce(
             dpnp.asnumpy(a), axis=axis, keepdims=keepdims, dtype=exp_dtype
         )
@@ -1938,11 +1948,17 @@ def test_reduce_hypot(self, dtype, axis, keepdims):
     @pytest.mark.parametrize("keepdims", [True, False])
     def test_reduce_hypot_out(self, dtype, axis, keepdims):
         a = dpnp.ones((3, 4, 5, 6, 7), dtype=dtype)
-        exp_dtype = dpnp.default_float_type(a.device)
+        exp_dtype = (
+            dpnp.default_float_type(a.device) if dtype == dpnp.bool else None
+        )
         exp = numpy.hypot.reduce(
             dpnp.asnumpy(a), axis=axis, keepdims=keepdims, dtype=exp_dtype
         )
-        dpnp_out = dpnp.empty(exp.shape, dtype=exp_dtype)
+
+        exp_dtype = exp.dtype
+        if exp_dtype == numpy.float64 and not has_support_aspect64():
+            exp_dtype = numpy.float32
+        dpnp_out = dpnp.empty_like(a, shape=exp.shape, dtype=exp_dtype)
         res = dpnp.reduce_hypot(a, axis=axis, out=dpnp_out, keepdims=keepdims)
 
         assert res is dpnp_out

diff --git a/tests/test_sum.py b/tests/test_sum.py
@@ -61,10 +61,7 @@ def test_sum_axis():
     ia = dpnp.array(a)
 
     result = dpnp.sum(ia, axis=1)
-    if has_support_aspect64():
-        expected = numpy.sum(a, axis=1, dtype=numpy.float64)
-    else:
-        expected = numpy.sum(a, axis=1)
+    expected = numpy.sum(a, axis=1)
     assert_array_equal(expected, result)