Update test_sum and test_prod to support complex dtypes

Author: asmeurer

array_api_tests/dtype_helpers.py

@@ -157,19 +157,18 @@ def is_int_dtype(dtype):
     return dtype in all_int_dtypes
 
 
-def is_float_dtype(dtype):
+def is_float_dtype(dtype, real=False):
     # None equals NumPy's xp.float64 object, so we specifically check it here.
     # xp.float64 is in fact an alias of np.dtype('float64'), and its equality
     # with None is meant to be deprecated at some point.
     # See https://github.com/numpy/numpy/issues/18434
     if dtype is None:
         return False
     valid_dtypes = real_float_dtypes
-    if api_version > "2021.12":
+    if api_version > "2021.12" and not real:
         valid_dtypes += complex_dtypes
     return dtype in valid_dtypes
 
-
 def get_scalar_type(dtype: DataType) -> ScalarType:
     if dtype in all_int_dtypes:
         return int

array_api_tests/test_statistical_functions.py

@@ -11,7 +11,7 @@
 from . import hypothesis_helpers as hh
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
-from . import xps
+from . import xps, api_version
 from ._array_module import _UndefinedStub
 from .typing import DataType
 
@@ -145,11 +145,19 @@ def test_prod(x, data):
                     _dtype = x.dtype
                 else:
                     _dtype = default_dtype
-        else:
+        elif dh.is_float_dtype(x.dtype, real=True):
             if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_float]:
                 _dtype = x.dtype
             else:
                 _dtype = dh.default_float
+        elif api_version > "2021.12":
+            # Complex dtype
+            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_complex]:
+                _dtype = x.dtype
+            else:
+                _dtype = dh.default_complex
+        else:
+            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
     else:
         _dtype = dtype
     if _dtype is None:
@@ -253,11 +261,19 @@ def test_sum(x, data):
                     _dtype = x.dtype
                 else:
                     _dtype = default_dtype
-        else:
+        elif dh.is_float_dtype(x.dtype, real=True):
             if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_float]:
                 _dtype = x.dtype
             else:
                 _dtype = dh.default_float
+        elif api_version > "2021.12":
+            # Complex dtype
+            if dh.dtype_nbits[x.dtype] > dh.dtype_nbits[dh.default_complex]:
+                _dtype = x.dtype
+            else:
+                _dtype = dh.default_complex
+        else:
+            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
     else:
         _dtype = dtype
     if _dtype is None: