From c26b752f3e125d4de2206f340ec8216332c32c82 Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Tue, 10 Nov 2026 18:01:21 +0100 Subject: [PATCH 1/7] complete tests with complex arrays --- tests/main/test_statistical.py | 34 ++++++++++++++++++++++++++++++++++ 1 file changed, 34 insertions(+) diff --git a/tests/main/test_statistical.py b/tests/main/test_statistical.py index a8fb1ea1..cc39e824 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 + 1 / 3j, 0.5 + 0.5j], + dtype=xp.complex128, + ) + assert_close(res, expected) From 334a7e2d69572d711ea43c603874f97cbaa0121b Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Tue, 10 Nov 2026 18:09:56 +0100 Subject: [PATCH 2/7] update conditions --- src/array_api_extra/_statistical.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) 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) From 56ffb642ada3f8cca8be900289737bd09ebeaeba Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Tue, 10 Nov 2026 18:19:00 +0100 Subject: [PATCH 3/7] correct expected complex value in nanmean --- tests/main/test_statistical.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/main/test_statistical.py b/tests/main/test_statistical.py index cc39e824..83cce564 100644 --- a/tests/main/test_statistical.py +++ b/tests/main/test_statistical.py @@ -600,7 +600,7 @@ def test_complex(self, xp: ArrayNamespace): res = nanmean(a, axis=0) expected = xp.asarray( - [1.5 + 0.5j, -0.5 + 2j, 10 / 3 + 1 / 3j, 0.5 + 0.5j], + [1.5 + 0.5j, -0.5 + 2j, 10 / 3 + 1j / 3, 0.5 + 0.5j], dtype=xp.complex128, ) assert_close(res, expected) From 64d89f14ab4ad60fcccb992b3976dacbe0571794 Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Wed, 11 Nov 2026 11:55:10 +0100 Subject: [PATCH 4/7] add private function --- src/array_api_extra/_statistical.py | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index 2dc7ef57..e6bc4862 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -321,6 +321,16 @@ def nanmax( return _agnostic._statistical.nanmax(a, axis=axis, xp=xp) +def _is_torch_and_realarray_or_notcpu( + a: Array, + xp: ArrayNamespace | None = None +) -> bool: + if not _compat.is_torch_namespace(xp): + return False + device = _compat.device(a) + return xp.isdtype(a.dtype, "real floating") or getattr(device, "type", None) != "cpu" + + def nanmean( a: Array, /, @@ -366,9 +376,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) and - (xp.isdtype(a.dtype, "real floating") or _compat.device(a).type != "cpu") - ) + or _is_torch_and_realarray_or_notcpu(a, xp) ): return xp.nanmean(a, axis=axis) @@ -420,9 +428,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) and - (xp.isdtype(a.dtype, "real floating") or _compat.device(a).type != "cpu") - ) + or _is_torch_and_realarray_or_notcpu(a, xp) ): return xp.nansum(a, axis=axis) From ea44f2cfe27b1ea5502c4e1a372a6b8c882c1f5b Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Wed, 11 Nov 2026 12:02:04 +0100 Subject: [PATCH 5/7] correct function --- src/array_api_extra/_statistical.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index e6bc4862..b34ffd81 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -323,7 +323,7 @@ def nanmax( def _is_torch_and_realarray_or_notcpu( a: Array, - xp: ArrayNamespace | None = None + xp: ArrayNamespace, ) -> bool: if not _compat.is_torch_namespace(xp): return False From 7d4239d8a4526a71c7b5c8601b650fc0936f85f3 Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Wed, 11 Nov 2026 12:14:39 +0100 Subject: [PATCH 6/7] correct again --- src/array_api_extra/_statistical.py | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index b34ffd81..80924542 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -325,10 +325,13 @@ def _is_torch_and_realarray_or_notcpu( a: Array, xp: ArrayNamespace, ) -> bool: + """Test if Torch namespace and if real array or not CPU""" if not _compat.is_torch_namespace(xp): return False - device = _compat.device(a) - return xp.isdtype(a.dtype, "real floating") or getattr(device, "type", None) != "cpu" + return ( + xp.isdtype(a.dtype, "real floating") + or getattr(_compat.device(a), "type", None) != "cpu" + ) def nanmean( From 13a161b6211cb26a0fef19258524103e1317d001 Mon Sep 17 00:00:00 2001 From: qbarthelemy Date: Thu, 12 Nov 2026 09:04:08 +0100 Subject: [PATCH 7/7] add doc --- src/array_api_extra/_statistical.py | 17 ++++++++++++++++- 1 file changed, 16 insertions(+), 1 deletion(-) diff --git a/src/array_api_extra/_statistical.py b/src/array_api_extra/_statistical.py index 80924542..3d180455 100644 --- a/src/array_api_extra/_statistical.py +++ b/src/array_api_extra/_statistical.py @@ -325,7 +325,22 @@ def _is_torch_and_realarray_or_notcpu( a: Array, xp: ArrayNamespace, ) -> bool: - """Test if Torch namespace and if real array or not CPU""" + """ + 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 (