Factor out dtype logic from test_sum() and test_prod() and apply it to test_trace()

Author: asmeurer

array_api_tests/dtype_helpers.py

@@ -242,6 +242,47 @@ def as_real_dtype(dtype):
     else:
         raise ValueError("as_real_dtype requires a floating-point dtype")
 
+def accumulation_result_dtype(x_dtype, dtype_kwarg):
+    """
+    Result dtype logic for sum(), prod(), and trace()
+
+    Note: may return None if a default uint cannot exist (e.g., for pytorch
+    which doesn't support uint32 or uint64). See https://github.com/data-apis/array-api-tests/issues/106
+
+    """
+    if dtype_kwarg is None:
+        if is_int_dtype(x_dtype):
+            if x_dtype in uint_dtypes:
+                default_dtype = default_uint
+            else:
+                default_dtype = default_int
+            if default_dtype is None:
+                _dtype = None
+            else:
+                m, M = dtype_ranges[x_dtype]
+                d_m, d_M = dtype_ranges[default_dtype]
+                if m < d_m or M > d_M:
+                    _dtype = x_dtype
+                else:
+                    _dtype = default_dtype
+        elif is_float_dtype(x_dtype, include_complex=False):
+            if dtype_nbits[x_dtype] > dtype_nbits[default_float]:
+                _dtype = x_dtype
+            else:
+                _dtype = default_float
+        elif api_version > "2021.12":
+            # Complex dtype
+            if dtype_nbits[x_dtype] > dtype_nbits[default_complex]:
+                _dtype = x_dtype
+            else:
+                _dtype = default_complex
+        else:
+            raise RuntimeError("Unexpected dtype. This indicates a bug in the test suite.")
+    else:
+        _dtype = dtype_kwarg
+
+    return _dtype
+
 if not hasattr(xp, "asarray"):
     default_int = xp.int32
     default_float = xp.float32

array_api_tests/test_linalg.py

@@ -795,11 +795,16 @@ def test_tensordot(x1, x2, kw):
 def test_trace(x, kw):
     res = linalg.trace(x, **kw)
 
-    # TODO: trace() should promote in some cases. See
-    # https://github.com/data-apis/array-api/issues/202. See also the dtype
-    # argument to sum() below.
-
-    # assert res.dtype == x.dtype, "trace() returned the wrong dtype"
+    dtype = kw.get("dtype", None)
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
+        # If a default uint cannot exist (i.e. in PyTorch which doesn't support
+        # uint32 or uint64), we skip testing the output dtype.
+        # See https://github.com/data-apis/array-api-tests/issues/160
+        if x.dtype in dh.uint_dtypes:
+            assert dh.is_int_dtype(res.dtype)  # sanity check
+    else:
+        ph.assert_dtype("trace", in_dtype=x.dtype, out_dtype=res.dtype, expected=expected_dtype)
 
     n, m = x.shape[-2:]
     ph.assert_result_shape('trace', x.shape, res.shape, expected=x.shape[:-2])

array_api_tests/test_statistical_functions.py

@@ -130,44 +130,15 @@ def test_prod(x, data):
         out = xp.prod(x, **kw)
 
     dtype = kw.get("dtype", None)
-    if dtype is None:
-        if dh.is_int_dtype(x.dtype):
-            if x.dtype in dh.uint_dtypes:
-                default_dtype = dh.default_uint
-            else:
-                default_dtype = dh.default_int
-            if default_dtype is None:
-                _dtype = None
-            else:
-                m, M = dh.dtype_ranges[x.dtype]
-                d_m, d_M = dh.dtype_ranges[default_dtype]
-                if m < d_m or M > d_M:
-                    _dtype = x.dtype
-                else:
-                    _dtype = default_dtype
-        elif dh.is_float_dtype(x.dtype, include_complex=False):
-            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:
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
         # If a default uint cannot exist (i.e. in PyTorch which doesn't support
         # uint32 or uint64), we skip testing the output dtype.
         # See https://github.com/data-apis/array-api-tests/issues/106
         if x.dtype in dh.uint_dtypes:
             assert dh.is_int_dtype(out.dtype)  # sanity check
     else:
-        ph.assert_dtype("prod", in_dtype=x.dtype, out_dtype=out.dtype, expected=_dtype)
+        ph.assert_dtype("prod", in_dtype=x.dtype, out_dtype=out.dtype, expected=expected_dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
         "prod", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw
@@ -246,44 +217,15 @@ def test_sum(x, data):
         out = xp.sum(x, **kw)
 
     dtype = kw.get("dtype", None)
-    if dtype is None:
-        if dh.is_int_dtype(x.dtype):
-            if x.dtype in dh.uint_dtypes:
-                default_dtype = dh.default_uint
-            else:
-                default_dtype = dh.default_int
-            if default_dtype is None:
-                _dtype = None
-            else:
-                m, M = dh.dtype_ranges[x.dtype]
-                d_m, d_M = dh.dtype_ranges[default_dtype]
-                if m < d_m or M > d_M:
-                    _dtype = x.dtype
-                else:
-                    _dtype = default_dtype
-        elif dh.is_float_dtype(x.dtype, include_complex=False):
-            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:
+    expected_dtype = dh.accumulation_result_dtype(x.dtype, dtype)
+    if expected_dtype is None:
         # If a default uint cannot exist (i.e. in PyTorch which doesn't support
         # uint32 or uint64), we skip testing the output dtype.
         # See https://github.com/data-apis/array-api-tests/issues/160
         if x.dtype in dh.uint_dtypes:
             assert dh.is_int_dtype(out.dtype)  # sanity check
     else:
-        ph.assert_dtype("sum", in_dtype=x.dtype, out_dtype=out.dtype, expected=_dtype)
+        ph.assert_dtype("sum", in_dtype=x.dtype, out_dtype=out.dtype, expected=expected_dtype)
     _axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
     ph.assert_keepdimable_shape(
         "sum", in_shape=x.shape, out_shape=out.shape, axes=_axes, keepdims=keepdims, kw=kw