Factor out dtype logic from test_sum() and test_prod() and apply it t…

…o test_trace()

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