Add more tests for dpnp.sum and sum_over_axis_0 extension (#1488)

* Add more tests for dpnp.sum and sum_over_axis_0 extension

* Add keepdims=True, bool and complex dtypes

---------

Co-authored-by: Anton <100830759+antonwolfy@users.noreply.github.com>

Author: AlexanderKalistratov

tests/test_extensions.py

@@ -193,16 +193,16 @@ def test_mean_over_axis_0_unsupported_out_types(
     input = dpt.empty((height, width), dtype=input_type, device=device)
     output = dpt.empty(width, dtype=output_type, device=device)
 
-    if func(input, output):
-        print(output_type)
     assert func(input, output) is None
 
 
 @pytest.mark.parametrize(
     "func, device, input_type, output_type",
     product(mean_sum, all_devices, [dpt.float32], [dpt.float32]),
 )
-def test_mean_over_axis_0_f_contig_input(func, device, input_type, output_type):
+def test_mean_sum_over_axis_0_f_contig_input(
+    func, device, input_type, output_type
+):
     skip_unsupported(device, input_type)
     skip_unsupported(device, output_type)
 
@@ -212,16 +212,14 @@ def test_mean_over_axis_0_f_contig_input(func, device, input_type, output_type):
     input = dpt.empty((height, width), dtype=input_type, device=device).T
     output = dpt.empty(width, dtype=output_type, device=device)
 
-    if func(input, output):
-        print(output_type)
     assert func(input, output) is None
 
 
 @pytest.mark.parametrize(
     "func, device, input_type, output_type",
     product(mean_sum, all_devices, [dpt.float32], [dpt.float32]),
 )
-def test_mean_over_axis_0_f_contig_output(
+def test_mean_sum_over_axis_0_f_contig_output(
     func, device, input_type, output_type
 ):
     skip_unsupported(device, input_type)
@@ -230,9 +228,25 @@ def test_mean_over_axis_0_f_contig_output(
     height = 1
     width = 10
 
-    input = dpt.empty((height, 10), dtype=input_type, device=device)
-    output = dpt.empty(20, dtype=output_type, device=device)[::2]
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width * 2, dtype=output_type, device=device)[::2]
+
+    assert func(input, output) is None
+
+
+@pytest.mark.parametrize(
+    "func, device, input_type, output_type",
+    product(mean_sum, all_devices, [dpt.float32], [dpt.float32, dpt.float64]),
+)
+def test_mean_sum_over_axis_0_big_output(func, device, input_type, output_type):
+    skip_unsupported(device, input_type)
+    skip_unsupported(device, output_type)
+
+    local_mem_size = device.local_mem_size
+    height = 1
+    width = 1 + local_mem_size // output_type.itemsize
+
+    input = dpt.empty((height, width), dtype=input_type, device=device)
+    output = dpt.empty(width, dtype=output_type, device=device)
 
-    if func(input, output):
-        print(output_type)
     assert func(input, output) is None

tests/test_mathematical.py

@@ -1,3 +1,5 @@
+from itertools import permutations
+
 import numpy
 import pytest
 from numpy.testing import (
@@ -1056,23 +1058,42 @@ def test_sum_empty_out(dtype):
 
 
 @pytest.mark.parametrize(
-    "shape", [(), (1, 2, 3), (1, 0, 2), (10), (3, 3, 3), (5, 5), (0, 6)]
-)
-@pytest.mark.parametrize(
-    "dtype_in", get_all_dtypes(no_complex=True, no_bool=True)
-)
-@pytest.mark.parametrize(
-    "dtype_out", get_all_dtypes(no_complex=True, no_bool=True)
+    "shape",
+    [
+        (),
+        (1, 2, 3),
+        (1, 0, 2),
+        (10,),
+        (3, 3, 3),
+        (5, 5),
+        (0, 6),
+        (10, 1),
+        (1, 10),
+    ],
 )
-def test_sum(shape, dtype_in, dtype_out):
-    a_np = numpy.ones(shape, dtype=dtype_in)
-    a = dpnp.ones(shape, dtype=dtype_in)
-    axes = [None, 0, 1, 2]
+@pytest.mark.parametrize("dtype_in", get_all_dtypes())
+@pytest.mark.parametrize("dtype_out", get_all_dtypes())
+@pytest.mark.parametrize("transpose", [True, False])
+@pytest.mark.parametrize("keepdims", [True, False])
+def test_sum(shape, dtype_in, dtype_out, transpose, keepdims):
+    size = numpy.prod(shape)
+    a_np = numpy.arange(size).astype(dtype_in).reshape(shape)
+    a = dpnp.asarray(a_np)
+
+    if transpose:
+        a_np = a_np.T
+        a = a.T
+
+    axes_range = list(numpy.arange(len(shape)))
+    axes = [None]
+    axes += axes_range
+    axes += permutations(axes_range, 2)
+    axes.append(tuple(axes_range))
+
     for axis in axes:
-        if axis is None or axis < a.ndim:
-            numpy_res = a_np.sum(axis=axis, dtype=dtype_out)
-            dpnp_res = a.sum(axis=axis, dtype=dtype_out)
-            assert_array_equal(numpy_res, dpnp_res.asnumpy())
+        numpy_res = a_np.sum(axis=axis, dtype=dtype_out, keepdims=keepdims)
+        dpnp_res = a.sum(axis=axis, dtype=dtype_out, keepdims=keepdims)
+        assert_array_equal(numpy_res, dpnp_res.asnumpy())
 
 
 class TestMean: