Remove mask usage in ph.assert_array()

Author: honno

array_api_tests/meta/test_pytest_helpers.py

@@ -1,7 +1,7 @@
 from pytest import raises
 
-from .. import pytest_helpers as ph
 from .. import _array_module as xp
+from .. import pytest_helpers as ph
 
 
 def test_assert_dtype():
@@ -11,3 +11,12 @@ def test_assert_dtype():
     ph.assert_dtype("bool_func", [xp.uint8, xp.int8], xp.bool, xp.bool)
     ph.assert_dtype("single_promoted_func", [xp.uint8], xp.uint8)
     ph.assert_dtype("single_bool_func", [xp.uint8], xp.bool, xp.bool)
+
+
+def test_assert_array():
+    ph.assert_array("int zeros", xp.asarray(0), xp.asarray(0))
+    ph.assert_array("pos zeros", xp.asarray(0.0), xp.asarray(0.0))
+    with raises(AssertionError):
+        ph.assert_array("mixed sign zeros", xp.asarray(0.0), xp.asarray(-0.0))
+    with raises(AssertionError):
+        ph.assert_array("mixed sign zeros", xp.asarray(-0.0), xp.asarray(0.0))

array_api_tests/pytest_helpers.py

@@ -24,6 +24,7 @@
     "assert_shape",
     "assert_result_shape",
     "assert_keepdimable_shape",
+    "assert_0d_equals",
     "assert_fill",
     "assert_array",
 ]
@@ -242,15 +243,28 @@ def assert_fill(
 def assert_array(func_name: str, out: Array, expected: Array, /, **kw):
     assert_dtype(func_name, out.dtype, expected.dtype)
     assert_shape(func_name, out.shape, expected.shape, **kw)
-    msg = f"out not as expected [{func_name}({fmt_kw(kw)})]\n{out=}\n{expected=}"
+    f_func = f"[{func_name}({fmt_kw(kw)})]"
     if dh.is_float_dtype(out.dtype):
-        neg_zeros = expected == -0.0
-        assert xp.all((out == -0.0) == neg_zeros), msg
-        pos_zeros = expected == +0.0
-        assert xp.all((out == +0.0) == pos_zeros), msg
-        nans = xp.isnan(expected)
-        assert xp.all(xp.isnan(out) == nans), msg
-        mask = ~(neg_zeros | pos_zeros | nans)
-        assert xp.all(out[mask] == expected[mask]), msg
+        for idx in sh.ndindex(out.shape):
+            at_out = out[idx]
+            at_expected = expected[idx]
+            msg = (
+                f"{sh.fmt_idx('out', idx)}={at_out}, should be {at_expected} "
+                f"{f_func}"
+            )
+            if xp.isnan(at_expected):
+                assert xp.isnan(at_out), msg
+            elif at_expected == 0.0 or at_expected == -0.0:
+                scalar_at_expected = float(at_expected)
+                scalar_at_out = float(at_out)
+                if is_pos_zero(scalar_at_expected):
+                    assert is_pos_zero(scalar_at_out), msg
+                else:
+                    assert is_neg_zero(scalar_at_expected)  # sanity check
+                    assert is_neg_zero(scalar_at_out), msg
+            else:
+                assert at_out == at_expected, msg
     else:
-        assert xp.all(out == expected), msg
+        assert xp.all(out == expected), (
+            f"out not as expected {f_func}\n" f"{out=}\n{expected=}"
+        )

array_api_tests/test_creation_functions.py

@@ -280,7 +280,7 @@ def test_asarray_arrays(x, data):
         if copy:
             assert not xp.all(
                 out == x
-            ), "xp.all(out == x)=True, but should be False after x was mutated\n{out=}"
+            ), f"xp.all(out == x)=True, but should be False after x was mutated\n{out=}"
         elif copy is False:
             pass  # TODO