Support "same sign except" special cases

See data-apis/array-api#92.

Author: asmeurer

array_api_tests/special_cases/test_multiply.py

@@ -31,28 +31,28 @@ def test_multiply_special_cases_two_args_either(arg1, arg2):
 
 
 @given(numeric_arrays, numeric_arrays)
-def test_multiply_special_cases_two_args_same_sign(arg1, arg2):
+def test_multiply_special_cases_two_args_same_sign_except(arg1, arg2):
     """
     Special case test for `multiply(x1, x2)`:
 
-        -   If `x1_i` and `x2_i` have the same mathematical sign, the result has a positive mathematical sign.
+        -   If `x1_i` and `x2_i` have the same mathematical sign, the result has a positive mathematical sign, except where it is `NaN` as in the rules below.
 
     """
     res = multiply(arg1, arg2)
-    mask = same_sign(arg1, arg2)
+    mask = logical_and(same_sign(arg1, arg2), logical_not(exactly_equal(res, NaN(res.shape, res.dtype))))
     assert_positive_mathematical_sign(res[mask])
 
 
 @given(numeric_arrays, numeric_arrays)
-def test_multiply_special_cases_two_args_different_signs(arg1, arg2):
+def test_multiply_special_cases_two_args_different_signs_except(arg1, arg2):
     """
     Special case test for `multiply(x1, x2)`:
 
-        -   If `x1_i` and `x2_i` have different mathematical signs, the result has a negative mathematical sign.
+        -   If `x1_i` and `x2_i` have different mathematical signs, the result has a negative mathematical sign, except where it is `NaN` as in the rules below.
 
     """
     res = multiply(arg1, arg2)
-    mask = logical_not(same_sign(arg1, arg2))
+    mask = logical_and(logical_not(same_sign(arg1, arg2)), logical_not(exactly_equal(res, NaN(res.shape, res.dtype))))
     assert_negative_mathematical_sign(res[mask])
 
 

generate_stubs.py

@@ -212,8 +212,10 @@ def {sig}:{doc}
     TWO_ARGS_EQUAL__EITHER = regex.compile(rf'^- +If `x1_i` is {_value} and `x2_i` is either {_value} or {_value}, the result is {_value}\.$'),
     TWO_ARGS_EITHER__EITHER = regex.compile(rf'^- +If `x1_i` is either {_value} or {_value} and `x2_i` is either {_value} or {_value}, the result is {_value}\.$'),
     TWO_ARGS_SAME_SIGN = regex.compile(rf'^- +If `x1_i` and `x2_i` have the same mathematical sign, the result has a {_value}\.$'),
+    TWO_ARGS_SAME_SIGN_EXCEPT = regex.compile(rf'^- +If `x1_i` and `x2_i` have the same mathematical sign, the result has a {_value}, except where it is {_value} as in the rules below\.$'),
     TWO_ARGS_SAME_SIGN_BOTH = regex.compile(rf'^- +If `x1_i` and `x2_i` have the same mathematical sign and are both {_value}, the result has a {_value}\.$'),
     TWO_ARGS_DIFFERENT_SIGNS = regex.compile(rf'^- +If `x1_i` and `x2_i` have different mathematical signs, the result has a {_value}\.$'),
+    TWO_ARGS_DIFFERENT_SIGNS_EXCEPT = regex.compile(rf'^- +If `x1_i` and `x2_i` have different mathematical signs, the result has a {_value}, except where it is {_value} as in the rules below\.$'),
     TWO_ARGS_DIFFERENT_SIGNS_BOTH = regex.compile(rf'^- +If `x1_i` and `x2_i` have different mathematical signs and are both {_value}, the result has a {_value}\.$'),
     TWO_ARGS_EVEN_IF = regex.compile(rf'^- +If `x2_i` is {_value}, the result is {_value}, even if `x1_i` is {_value}\.$'),
 
@@ -483,6 +485,11 @@ def generate_special_case_test(func, typ, m, test_name_extra, sigs):
             result, = m.groups()
             mask = "same_sign(arg1, arg2)"
             assertion = get_assert("exactly_equal", result)
+        elif typ == "TWO_ARGS_SAME_SIGN_EXCEPT":
+            result, value = m.groups()
+            value = parse_value(value, "res")
+            mask = f"logical_and(same_sign(arg1, arg2), logical_not(exactly_equal(res, {value})))"
+            assertion = get_assert("exactly_equal", result)
         elif typ == "TWO_ARGS_SAME_SIGN_BOTH":
             value, result = m.groups()
             mask1 = get_mask("exactly_equal", "arg1", value)
@@ -493,6 +500,11 @@ def generate_special_case_test(func, typ, m, test_name_extra, sigs):
             result, = m.groups()
             mask = "logical_not(same_sign(arg1, arg2))"
             assertion = get_assert("exactly_equal", result)
+        elif typ == "TWO_ARGS_DIFFERENT_SIGNS_EXCEPT":
+            result, value = m.groups()
+            value = parse_value(value, "res")
+            mask = f"logical_and(logical_not(same_sign(arg1, arg2)), logical_not(exactly_equal(res, {value})))"
+            assertion = get_assert("exactly_equal", result)
         elif typ == "TWO_ARGS_DIFFERENT_SIGNS_BOTH":
             value, result = m.groups()
             mask1 = get_mask("exactly_equal", "arg1", value)