Make sure both arguments in the assertion are masked

Author: asmeurer

array_api_tests/special_cases/test_add.py

@@ -207,7 +207,7 @@ def test_add_special_cases_two_args_either__equal(arg1, arg2):
     """
     res = add(arg1, arg2)
     mask = logical_and(logical_or(exactly_equal(arg1, zero(arg1.dtype)), exactly_equal(arg1, -zero(arg1.dtype))), logical_and(isfinite(arg2), nonzero(arg2)))
-    assert_exactly_equal(res[mask], arg2)
+    assert_exactly_equal(res[mask], arg2[mask])
 
 
 @given(numeric_arrays, numeric_arrays)
@@ -220,7 +220,7 @@ def test_add_special_cases_two_args_equal__either(arg1, arg2):
     """
     res = add(arg1, arg2)
     mask = logical_and(logical_and(isfinite(arg1), nonzero(arg1)), logical_or(exactly_equal(arg2, zero(arg2.dtype)), exactly_equal(arg2, -zero(arg2.dtype))))
-    assert_exactly_equal(res[mask], arg1)
+    assert_exactly_equal(res[mask], arg1[mask])
 
 # TODO: Implement REMAINING test for:
 # -   In the remaining cases, when neither `infinity`, `+0`, `-0`, nor a `NaN` is involved, and the operands have the same mathematical sign or have different magnitudes, the sum must be computed and rounded to the nearest representable value according to IEEE 754-2019 and a supported round mode. If the magnitude is too large to represent, the operation overflows and the result is an `infinity` of appropriate mathematical sign.

array_api_tests/special_cases/test_ceil.py

@@ -24,4 +24,4 @@ def test_ceil_special_cases_one_arg_equal(arg1):
     """
     res = ceil(arg1)
     mask = isintegral(arg1)
-    assert_exactly_equal(res[mask], arg1)
+    assert_exactly_equal(res[mask], arg1[mask])

array_api_tests/special_cases/test_floor.py

@@ -24,4 +24,4 @@ def test_floor_special_cases_one_arg_equal(arg1):
     """
     res = floor(arg1)
     mask = isintegral(arg1)
-    assert_exactly_equal(res[mask], arg1)
+    assert_exactly_equal(res[mask], arg1[mask])

array_api_tests/special_cases/test_round.py

@@ -25,7 +25,7 @@ def test_round_special_cases_one_arg_equal(arg1):
     """
     res = round(arg1)
     mask = isintegral(arg1)
-    assert_exactly_equal(res[mask], arg1)
+    assert_exactly_equal(res[mask], arg1[mask])
 
 
 @given(numeric_arrays)

array_api_tests/special_cases/test_trunc.py

@@ -24,4 +24,4 @@ def test_trunc_special_cases_one_arg_equal(arg1):
     """
     res = trunc(arg1)
     mask = isintegral(arg1)
-    assert_exactly_equal(res[mask], arg1)
+    assert_exactly_equal(res[mask], arg1[mask])

generate_stubs.py

@@ -315,11 +315,11 @@ def get_assert(typ, result):
         _check_exactly_equal(typ, result)
         return "assert_negative(res[mask])"
     elif 'x_i' in result:
-        return f"assert_{typ}(res[mask], {result.replace('x_i', 'arg1')})"
+        return f"assert_{typ}(res[mask], {result.replace('x_i', 'arg1')}[mask])"
     elif 'x1_i' in result:
-        return f"assert_{typ}(res[mask], {result.replace('x1_i', 'arg1')})"
+        return f"assert_{typ}(res[mask], {result.replace('x1_i', 'arg1')}[mask])"
     elif 'x2_i' in result:
-        return f"assert_{typ}(res[mask], {result.replace('x2_i', 'arg2')})"
+        return f"assert_{typ}(res[mask], {result.replace('x2_i', 'arg2')}[mask])"
     # TODO: Get use something better than arg1 here for the arg
     result = parse_value(result, "arg1")
     try: