Test mean()

Author: honno

array_api_tests/test_statistical_functions.py

@@ -10,6 +10,8 @@
 from . import pytest_helpers as ph
 from . import xps
 
+RTOL = 0.05
+
 
 @given(
     x=xps.arrays(dtype=xps.numeric_dtypes(), shape=hh.shapes(min_side=1)),
@@ -37,7 +39,7 @@ def test_min(x, data):
         if keepdims:
             idx = tuple(1 for _ in x.shape)
             msg = f"{out.shape=}, should be reduced dimension {idx} [{f_func}]"
-            assert out.shape == idx
+            assert out.shape == idx, msg
         else:
             ph.assert_shape("min", out.shape, (), **kw)
 
@@ -84,7 +86,7 @@ def test_max(x, data):
         if keepdims:
             idx = tuple(1 for _ in x.shape)
             msg = f"{out.shape=}, should be reduced dimension {idx} [{f_func}]"
-            assert out.shape == idx
+            assert out.shape == idx, msg
         else:
             ph.assert_shape("max", out.shape, (), **kw)
 
@@ -105,11 +107,47 @@ def test_max(x, data):
                 assert max_ == expected, msg
 
 
-# TODO: generate kwargs
-@given(xps.arrays(dtype=xps.floating_dtypes(), shape=hh.shapes(min_side=1)))
-def test_mean(x):
-    xp.mean(x)
-    # TODO
+@given(
+    x=xps.arrays(dtype=xps.floating_dtypes(), shape=hh.shapes(min_side=1)),
+    data=st.data(),
+)
+def test_mean(x, data):
+    axis_strats = [st.none()]
+    if x.shape != ():
+        axis_strats.append(
+            st.integers(-x.ndim, x.ndim - 1) | xps.valid_tuple_axes(x.ndim)
+        )
+    kw = data.draw(
+        hh.kwargs(axis=st.one_of(axis_strats), keepdims=st.booleans()), label="kw"
+    )
+
+    out = xp.mean(x, **kw)
+
+    ph.assert_dtype("mean", x.dtype, out.dtype)
+
+    f_func = f"mean({ph.fmt_kw(kw)})"
+
+    # TODO: support axis
+    if kw.get("axis") is None:
+        keepdims = kw.get("keepdims", False)
+        if keepdims:
+            idx = tuple(1 for _ in x.shape)
+            msg = f"{out.shape=}, should be reduced dimension {idx} [{f_func}]"
+            assert out.shape == idx, msg
+        else:
+            ph.assert_shape("max", out.shape, (), **kw)
+
+        # TODO: figure out NaN behaviour
+        if not xp.any(xp.isnan(x)):
+            _out = xp.reshape(out, ()) if keepdims else out
+            elements = []
+            for idx in ah.ndindex(x.shape):
+                s = float(x[idx])
+                elements.append(s)
+            mean = float(_out)
+            expected = sum(elements) / len(elements)
+            msg = f"out={mean}, should be roughly {expected} [{f_func}]"
+            assert math.isclose(mean, expected, rel_tol=RTOL), msg
 
 
 # TODO: generate kwargs