Add support for complex floating-point data types in mean (#796)

Author: tomwhite

cubed/array_api/statistical_functions.py

@@ -1,6 +1,8 @@
 import math
 
+from cubed.array_api.data_type_functions import isdtype
 from cubed.array_api.dtypes import (
+    _floating_dtypes,
     _real_floating_dtypes,
     _real_numeric_dtypes,
     _upcast_integral_dtypes,
@@ -94,14 +96,17 @@ def max(x, /, *, axis=None, keepdims=False, split_every=None):
 
 
 def mean(x, /, *, axis=None, keepdims=False, split_every=None):
-    if x.dtype not in _real_floating_dtypes:
-        raise TypeError("Only real floating-point dtypes are allowed in mean")
+    if x.dtype not in _floating_dtypes:
+        raise TypeError("Only floating-point dtypes are allowed in mean")
     # This implementation uses a Zarr group of two arrays to store a
     # pair of fields needed to keep per-chunk counts and totals for computing
     # the mean.
     dtype = x.dtype
     # TODO(#658): Should these be default dtypes?
-    intermediate_dtype = [("n", nxp.int64), ("total", nxp.float64)]
+    if isdtype(x.dtype, "complex floating"):
+        intermediate_dtype = [("n", nxp.int64), ("total", nxp.complex128)]
+    else:
+        intermediate_dtype = [("n", nxp.int64), ("total", nxp.float64)]
     extra_func_kwargs = dict(dtype=intermediate_dtype)
     return reduction(
         x,

cubed/tests/test_array_api.py

@@ -904,6 +904,11 @@ def test_mean_axis_0(spec, executor):
         np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0], [7.0, 8.0, 9.0]]).mean(axis=0),
     )
 
+def test_mean_complex():
+    a = xp.asarray([1.0+1.0j, 2.0+2.0j, 3.0+3.0j], chunks=(2,))
+    b = xp.mean(a)
+    assert_array_equal(b.compute(), np.array([1.0+1.0j, 2.0+2.0j, 3.0+3.0j]).mean())
+
 
 def test_sum(spec, executor):
     a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2), spec=spec)