Move assertion helpers to pytest_helpers.py

Author: honno

array_api_tests/pytest_helpers.py

@@ -1,12 +1,13 @@
-from .algos import broadcast_shapes
 import math
 from inspect import getfullargspec
 from typing import Any, Dict, Optional, Tuple, Union
 
+from . import _array_module as xp
 from . import array_helpers as ah
 from . import dtype_helpers as dh
 from . import function_stubs
-from .typing import Array, DataType, Scalar, Shape
+from .algos import broadcast_shapes
+from .typing import Array, DataType, Scalar, ScalarType, Shape
 
 __all__ = [
     "raises",
@@ -17,8 +18,10 @@
     "assert_kw_dtype",
     "assert_default_float",
     "assert_default_int",
+    "assert_default_index",
     "assert_shape",
     "assert_result_shape",
+    "assert_keepdimable_shape",
     "assert_fill",
 ]
 
@@ -117,6 +120,15 @@ def assert_default_int(func_name: str, dtype: DataType):
     assert dtype == dh.default_int, msg
 
 
+def assert_default_index(func_name: str, dtype: DataType, repr_name="out.dtype"):
+    f_dtype = dh.dtype_to_name[dtype]
+    msg = (
+        f"{repr_name}={f_dtype}, should be the default index dtype, "
+        f"which is either int32 or int64 [{func_name}()]"
+    )
+    assert dtype in (xp.int32, xp.int64), msg
+
+
 def assert_shape(
     func_name: str,
     out_shape: Union[int, Shape],
@@ -155,6 +167,57 @@ def assert_result_shape(
     assert out_shape == expected, msg
 
 
+def assert_keepdimable_shape(
+    func_name: str,
+    out_shape: Shape,
+    in_shape: Shape,
+    axes: Tuple[int, ...],
+    keepdims: bool,
+    /,
+    **kw,
+):
+    if keepdims:
+        shape = tuple(1 if axis in axes else side for axis, side in enumerate(in_shape))
+    else:
+        shape = tuple(side for axis, side in enumerate(in_shape) if axis not in axes)
+    assert_shape(func_name, out_shape, shape, **kw)
+
+
+def assert_0d_equals(
+    func_name: str, x_repr: str, x_val: Array, out_repr: str, out_val: Array, **kw
+):
+    msg = (
+        f"{out_repr}={out_val}, should be {x_repr}={x_val} "
+        f"[{func_name}({fmt_kw(kw)})]"
+    )
+    if dh.is_float_dtype(out_val.dtype) and xp.isnan(out_val):
+        assert xp.isnan(x_val), msg
+    else:
+        assert x_val == out_val, msg
+
+
+def assert_scalar_equals(
+    func_name: str,
+    type_: ScalarType,
+    idx: Shape,
+    out: Scalar,
+    expected: Scalar,
+    /,
+    **kw,
+):
+    out_repr = "out" if idx == () else f"out[{idx}]"
+    f_func = f"{func_name}({fmt_kw(kw)})"
+    if type_ is bool or type_ is int:
+        msg = f"{out_repr}={out}, should be {expected} [{f_func}]"
+        assert out == expected, msg
+    elif math.isnan(expected):
+        msg = f"{out_repr}={out}, should be {expected} [{f_func}]"
+        assert math.isnan(out), msg
+    else:
+        msg = f"{out_repr}={out}, should be roughly {expected} [{f_func}]"
+        assert math.isclose(out, expected, rel_tol=0.25, abs_tol=1), msg
+
+
 def assert_fill(
     func_name: str, fill_value: Scalar, dtype: DataType, out: Array, /, **kw
 ):

array_api_tests/test_manipulation_functions.py

@@ -44,19 +44,6 @@ def assert_array_ndindex(
             assert out[out_idx] == x[x_idx], msg
 
 
-def assert_equals(
-    func_name: str, x_repr: str, x_val: Array, out_repr: str, out_val: Array, **kw
-):
-    msg = (
-        f"{out_repr}={out_val}, should be {x_repr}={x_val} "
-        f"[{func_name}({ph.fmt_kw(kw)})]"
-    )
-    if dh.is_float_dtype(out_val.dtype) and xp.isnan(out_val):
-        assert xp.isnan(x_val), msg
-    else:
-        assert x_val == out_val, msg
-
-
 @st.composite
 def concat_shapes(draw, shape, axis):
     shape = list(shape)
@@ -104,7 +91,7 @@ def test_concat(dtypes, kw, data):
         for x_num, x in enumerate(arrays, 1):
             for x_idx in sh.ndindex(x.shape):
                 out_i = next(out_indices)
-                assert_equals(
+                ph.assert_0d_equals(
                     "concat",
                     f"x{x_num}[{x_idx}]",
                     x[x_idx],
@@ -120,7 +107,7 @@ def test_concat(dtypes, kw, data):
                 indexed_x = x[idx]
                 for x_idx in sh.ndindex(indexed_x.shape):
                     out_idx = next(out_indices)
-                    assert_equals(
+                    ph.assert_0d_equals(
                         "concat",
                         f"x{x_num}[{f_idx}][{x_idx}]",
                         indexed_x[x_idx],
@@ -360,7 +347,7 @@ def test_stack(shape, dtypes, kw, data):
             indexed_x = x[idx]
             for x_idx in sh.ndindex(indexed_x.shape):
                 out_idx = next(out_indices)
-                assert_equals(
+                ph.assert_0d_equals(
                     "stack",
                     f"x{x_num}[{f_idx}][{x_idx}]",
                     indexed_x[x_idx],

array_api_tests/test_searching_functions.py

@@ -8,18 +8,6 @@
 from . import shape_helpers as sh
 from . import xps
 from .algos import broadcast_shapes
-from .test_manipulation_functions import assert_equals as assert_equals_
-from .test_statistical_functions import assert_equals, assert_keepdimable_shape
-from .typing import DataType
-
-
-def assert_default_index(func_name: str, dtype: DataType, repr_name="out.dtype"):
-    f_dtype = dh.dtype_to_name[dtype]
-    msg = (
-        f"{repr_name}={f_dtype}, should be the default index dtype, "
-        f"which is either int32 or int64 [{func_name}()]"
-    )
-    assert dtype in (xp.int32, xp.int64), msg
 
 
 @given(
@@ -41,9 +29,9 @@ def test_argmax(x, data):
 
     out = xp.argmax(x, **kw)
 
-    assert_default_index("argmax", out.dtype)
+    ph.assert_default_index("argmax", out.dtype)
     axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
-    assert_keepdimable_shape(
+    ph.assert_keepdimable_shape(
         "argmax", out.shape, x.shape, axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
@@ -54,7 +42,7 @@ def test_argmax(x, data):
             s = scalar_type(x[idx])
             elements.append(s)
         expected = max(range(len(elements)), key=elements.__getitem__)
-        assert_equals("argmax", int, out_idx, max_i, expected)
+        ph.assert_scalar_equals("argmax", int, out_idx, max_i, expected)
 
 
 @given(
@@ -76,9 +64,9 @@ def test_argmin(x, data):
 
     out = xp.argmin(x, **kw)
 
-    assert_default_index("argmin", out.dtype)
+    ph.assert_default_index("argmin", out.dtype)
     axes = sh.normalise_axis(kw.get("axis", None), x.ndim)
-    assert_keepdimable_shape(
+    ph.assert_keepdimable_shape(
         "argmin", out.shape, x.shape, axes, kw.get("keepdims", False), **kw
     )
     scalar_type = dh.get_scalar_type(x.dtype)
@@ -89,7 +77,7 @@ def test_argmin(x, data):
             s = scalar_type(x[idx])
             elements.append(s)
         expected = min(range(len(elements)), key=elements.__getitem__)
-        assert_equals("argmin", int, out_idx, min_i, expected)
+        ph.assert_scalar_equals("argmin", int, out_idx, min_i, expected)
 
 
 # TODO: skip if opted out
@@ -106,7 +94,7 @@ def test_nonzero(x):
         assert (
             out[i].size == size
         ), f"out[{i}].size={x.size}, but should be out[0].size={size}"
-        assert_default_index("nonzero", out[i].dtype, repr_name=f"out[{i}].dtype")
+        ph.assert_default_index("nonzero", out[i].dtype, repr_name=f"out[{i}].dtype")
     indices = []
     if x.dtype == xp.bool:
         for idx in sh.ndindex(x.shape):
@@ -151,6 +139,10 @@ def test_where(shapes, dtypes, data):
     _x2 = xp.broadcast_to(x2, shape)
     for idx in sh.ndindex(shape):
         if _cond[idx]:
-            assert_equals_("where", f"_x1[{idx}]", _x1[idx], f"out[{idx}]", out[idx])
+            ph.assert_0d_equals(
+                "where", f"_x1[{idx}]", _x1[idx], f"out[{idx}]", out[idx]
+            )
         else:
-            assert_equals_("where", f"_x2[{idx}]", _x2[idx], f"out[{idx}]", out[idx])
+            ph.assert_0d_equals(
+                "where", f"_x2[{idx}]", _x2[idx], f"out[{idx}]", out[idx]
+            )

array_api_tests/test_set_functions.py

@@ -10,7 +10,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .test_searching_functions import assert_default_index
 
 
 @given(
@@ -29,11 +28,15 @@ def test_unique_all(x):
     ph.assert_dtype(
         "unique_all", x.dtype, out.values.dtype, repr_name="out.values.dtype"
     )
-    assert_default_index("unique_all", out.indices.dtype, repr_name="out.indices.dtype")
-    assert_default_index(
+    ph.assert_default_index(
+        "unique_all", out.indices.dtype, repr_name="out.indices.dtype"
+    )
+    ph.assert_default_index(
         "unique_all", out.inverse_indices.dtype, repr_name="out.inverse_indices.dtype"
     )
-    assert_default_index("unique_all", out.counts.dtype, repr_name="out.counts.dtype")
+    ph.assert_default_index(
+        "unique_all", out.counts.dtype, repr_name="out.counts.dtype"
+    )
 
     assert (
         out.indices.shape == out.values.shape
@@ -121,7 +124,7 @@ def test_unique_counts(x):
     ph.assert_dtype(
         "unique_counts", x.dtype, out.values.dtype, repr_name="out.values.dtype"
     )
-    assert_default_index(
+    ph.assert_default_index(
         "unique_counts", out.counts.dtype, repr_name="out.counts.dtype"
     )
     assert (
@@ -168,7 +171,7 @@ def test_unique_inverse(x):
     ph.assert_dtype(
         "unique_inverse", x.dtype, out.values.dtype, repr_name="out.values.dtype"
     )
-    assert_default_index(
+    ph.assert_default_index(
         "unique_inverse",
         out.inverse_indices.dtype,
         repr_name="out.inverse_indices.dtype",

array_api_tests/test_sorting.py

@@ -8,9 +8,6 @@
 from . import pytest_helpers as ph
 from . import shape_helpers as sh
 from . import xps
-from .test_manipulation_functions import assert_equals as assert_equals_
-from .test_searching_functions import assert_default_index
-from .test_statistical_functions import assert_equals
 
 
 # TODO: Test with signed zeros and NaNs (and ignore them somehow)
@@ -37,7 +34,7 @@ def test_argsort(x, data):
 
     out = xp.argsort(x, **kw)
 
-    assert_default_index("sort", out.dtype)
+    ph.assert_default_index("sort", out.dtype)
     ph.assert_shape("sort", out.shape, x.shape, **kw)
     axis = kw.get("axis", -1)
     axes = sh.normalise_axis(axis, x.ndim)
@@ -50,7 +47,7 @@ def test_argsort(x, data):
             # sorted(..., reverse=descending) doesn't always work
             indices_order = reversed(indices_order)
         for idx, o in zip(indices, indices_order):
-            assert_equals("argsort", int, idx, int(out[idx]), o)
+            ph.assert_scalar_equals("argsort", int, idx, int(out[idx]), o)
 
 
 # TODO: Test with signed zeros and NaNs (and ignore them somehow)
@@ -90,7 +87,7 @@ def test_sort(x, data):
         )
         x_indices = [indices[o] for o in indices_order]
         for out_idx, x_idx in zip(indices, x_indices):
-            assert_equals_(
+            ph.assert_0d_equals(
                 "sort",
                 f"x[{x_idx}]",
                 x[x_idx],