Fix numpy/cupy sum(), prod(), and trace() for complex upcasting with dtype=None

Author: asmeurer

array_api_compat/common/_aliases.py

@@ -396,9 +396,12 @@ def sum(
     keepdims: bool = False,
     **kwargs,
 ) -> ndarray:
-    # `xp.sum` already upcasts integers, but not floats
-    if dtype is None and x.dtype == xp.float32:
-        dtype = xp.float64
+    # `xp.sum` already upcasts integers, but not floats or complexes
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
     return xp.sum(x, axis=axis, dtype=dtype, keepdims=keepdims, **kwargs)
 
 def prod(
@@ -411,8 +414,11 @@ def prod(
     keepdims: bool = False,
     **kwargs,
 ) -> ndarray:
-    if dtype is None and x.dtype == xp.float32:
-        dtype = xp.float64
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
     return xp.prod(x, dtype=dtype, axis=axis, keepdims=keepdims, **kwargs)
 
 # ceil, floor, and trunc return integers for integer inputs

array_api_compat/common/_linalg.py

@@ -136,8 +136,13 @@ def vector_norm(x: ndarray, /, xp, *, axis: Optional[Union[int, Tuple[int, ...]]
 def diagonal(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray:
     return xp.diagonal(x, offset=offset, axis1=-2, axis2=-1, **kwargs)
 
-def trace(x: ndarray, /, xp, *, offset: int = 0, **kwargs) -> ndarray:
-    return xp.asarray(xp.trace(x, offset=offset, axis1=-2, axis2=-1, **kwargs))
+def trace(x: ndarray, /, xp, *, offset: int = 0, dtype=None, **kwargs) -> ndarray:
+    if dtype is None:
+        if x.dtype == xp.float32:
+            dtype = xp.float64
+        elif x.dtype == xp.complex64:
+            dtype = xp.complex128
+    return xp.asarray(xp.trace(x, offset=offset, dtype=dtype, axis1=-2, axis2=-1, **kwargs))
 
 __all__ = ['cross', 'matmul', 'outer', 'tensordot', 'EighResult',
            'QRResult', 'SlogdetResult', 'SVDResult', 'eigh', 'qr', 'slogdet',