diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index f9b6f6f2..2dc7ef57 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -366,7 +366,9 @@ 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 (_compat.is_torch_namespace(xp) and + (xp.isdtype(a.dtype, "real floating") or _compat.device(a).type != "cpu") + ) ): return xp.nanmean(a, axis=axis) @@ -418,7 +420,9 @@ 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 (_compat.is_torch_namespace(xp) and + (xp.isdtype(a.dtype, "real floating") or _compat.device(a).type != "cpu") + ) ): return xp.nansum(a, axis=axis) diff --git a/tests/main/test_statistical.py b/tests/main/test_statistical.py index a8fb1ea1..83cce564 100644 --- a/tests/main/test_statistical.py +++ b/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)