ENH: enhance nansum and nanmean to support complex arrays with torch on cpu

src/array_api_extra/_statistical.py

@@ -321,6 +321,34 @@ def nanmax(
     return _agnostic._statistical.nanmax(a, axis=axis, xp=xp)
 
 
+def _is_torch_and_realarray_or_notcpu(
+    a: Array,
+    xp: ArrayNamespace,
+) -> bool:
+    """
+    Return whether `xp` is Torch and `a` is real or not on the CPU.
+
+    Parameters
+    ----------
+    a : Array
+        Input array.
+    xp : ArrayNamespace
+        Namespace used to process `a`.
+
+    Returns
+    -------
+    bool
+        Whether `xp` is Torch and either `a` has a real floating-point dtype
+        or is located on a non-CPU device.
+    """
+    if not _compat.is_torch_namespace(xp):
+        return False
+    return (
+        xp.isdtype(a.dtype, "real floating")
+        or getattr(_compat.device(a), "type", None) != "cpu"
+    )
+
+
 def nanmean(
     a: Array,
     /,
@@ -366,7 +394,7 @@ def nanmean(
         or _compat.is_cupy_namespace(xp)
         or _compat.is_dask_namespace(xp)
         or _compat.is_jax_namespace(xp)
-        or _compat.is_torch_namespace(xp)
+        or _is_torch_and_realarray_or_notcpu(a, xp)
     ):
         return xp.nanmean(a, axis=axis)
 
@@ -418,7 +446,7 @@ def nansum(
         or _compat.is_cupy_namespace(xp)
         or _compat.is_dask_namespace(xp)
         or _compat.is_jax_namespace(xp)
-        or _compat.is_torch_namespace(xp)
+        or _is_torch_and_realarray_or_notcpu(a, xp)
     ):
         return xp.nansum(a, axis=axis)
 

tests/main/test_statistical.py

@@ -508,6 +508,23 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespa
         expected = xp.asarray(expected_list)
         assert_equal(res, expected)
 
+    def test_complex(self, xp: ArrayNamespace):
+        a = xp.asarray(
+            [
+                [1 + 2j, complex(float("nan"), 0.0), 4 + 1j, 5 - 1j],
+                [complex(float("nan"), 0.0), -2 + 1j, 3 - 2j, -4 + 2j],
+                [2 - 1j, 1 + 3j, 3 + 2j, complex(float("nan"), 0.0)],
+            ],
+            dtype=xp.complex128,
+        )
+
+        res = nansum(a, axis=0)
+        expected = xp.asarray(
+            [3 + 1j, -1 + 4j, 10 + 1j, 1 + 1j],
+            dtype=xp.complex128,
+        )
+        assert_equal(res, expected)
+
 
 class TestNanMean:
     def test_simple(self, xp: ArrayNamespace):
@@ -570,3 +587,20 @@ def test_xp(self, axis: int | None, expected_list: list[float], xp: ArrayNamespa
         res = nanmean(a, axis=axis, xp=xp)
         expected = xp.asarray(expected_list)
         assert_equal(res, expected)
+
+    def test_complex(self, xp: ArrayNamespace):
+        a = xp.asarray(
+            [
+                [1 + 2j, complex(float("nan"), 0.0), 4 + 1j, 5 - 1j],
+                [complex(float("nan"), 0.0), -2 + 1j, 3 - 2j, -4 + 2j],
+                [2 - 1j, 1 + 3j, 3 + 2j, complex(float("nan"), 0.0)],
+            ],
+            dtype=xp.complex128,
+        )
+
+        res = nanmean(a, axis=0)
+        expected = xp.asarray(
+            [1.5 + 0.5j, -0.5 + 2j, 10 / 3 + 1j / 3, 0.5 + 0.5j],
+            dtype=xp.complex128,
+        )
+        assert_close(res, expected)