Fix the special cases for logaddexp

data-apis/array-api#239

Author: asmeurer

array_api_tests/special_cases/test_logaddexp.py

@@ -7,19 +7,20 @@
 not modify it directly.
 """
 
-from ..array_helpers import NaN, assert_exactly_equal, exactly_equal, infinity, logical_or
+from ..array_helpers import (NaN, assert_exactly_equal, exactly_equal, infinity, logical_and,
+                             logical_not, logical_or)
 from ..hypothesis_helpers import numeric_arrays
 from .._array_module import logaddexp
 
 from hypothesis import given
 
 
 @given(numeric_arrays, numeric_arrays)
-def test_logaddexp_special_cases_two_args_either_1(arg1, arg2):
+def test_logaddexp_special_cases_two_args_either(arg1, arg2):
     """
     Special case test for `logaddexp(x1, x2)`:
 
-        - If either `x1_i` or `x2_i` is `NaN`, the result is `NaN`.
+        -   If either `x1_i` or `x2_i` is `NaN`, the result is `NaN`.
 
     """
     res = logaddexp(arg1, arg2)
@@ -28,13 +29,26 @@ def test_logaddexp_special_cases_two_args_either_1(arg1, arg2):
 
 
 @given(numeric_arrays, numeric_arrays)
-def test_logaddexp_special_cases_two_args_either_2(arg1, arg2):
+def test_logaddexp_special_cases_two_args_equal__notequal(arg1, arg2):
     """
     Special case test for `logaddexp(x1, x2)`:
 
-        - If either `x1_i` or `x2_i` is `+infinity`, the result is `+infinity`.
+        -   If `x1_i` is `+infinity` and `x2_i` is not `NaN`, the result is `+infinity`.
 
     """
     res = logaddexp(arg1, arg2)
-    mask = logical_or(exactly_equal(arg1, infinity(arg1.shape, arg1.dtype)), exactly_equal(arg2, infinity(arg1.shape, arg1.dtype)))
+    mask = logical_and(exactly_equal(arg1, infinity(arg1.shape, arg1.dtype)), logical_not(exactly_equal(arg2, NaN(arg2.shape, arg2.dtype))))
+    assert_exactly_equal(res[mask], (infinity(arg1.shape, arg1.dtype))[mask])
+
+
+@given(numeric_arrays, numeric_arrays)
+def test_logaddexp_special_cases_two_args_notequal__equal(arg1, arg2):
+    """
+    Special case test for `logaddexp(x1, x2)`:
+
+        -   If `x1_i` is not `NaN` and `x2_i` is `+infinity`, the result is `+infinity`.
+
+    """
+    res = logaddexp(arg1, arg2)
+    mask = logical_and(logical_not(exactly_equal(arg1, NaN(arg1.shape, arg1.dtype))), exactly_equal(arg2, infinity(arg2.shape, arg2.dtype)))
     assert_exactly_equal(res[mask], (infinity(arg1.shape, arg1.dtype))[mask])

generate_stubs.py

@@ -317,7 +317,7 @@ def {annotated_sig}:{doc}
     TWO_ARGS_EQUAL__LESS_NOTEQUAL = regex.compile(rf'^- +If `x1_i` is {_value}, `x2_i` is less than {_value}, and `x2_i` is not {_value}, the result is {_value}\.$'),
     TWO_ARGS_EQUAL__GREATER_EQUAL = regex.compile(rf'^- +If `x1_i` is {_value}, `x2_i` is greater than {_value}, and `x2_i` is {_value}, the result is {_value}\.$'),
     TWO_ARGS_EQUAL__GREATER_NOTEQUAL = regex.compile(rf'^- +If `x1_i` is {_value}, `x2_i` is greater than {_value}, and `x2_i` is not {_value}, the result is {_value}\.$'),
-    TWO_ARGS_NOTEQUAL__EQUAL = regex.compile(rf'^- +If `x1_i` is not equal to {_value} and `x2_i` is {_value}, the result is {_value}\.$'),
+    TWO_ARGS_NOTEQUAL__EQUAL = regex.compile(rf'^- +If `x1_i` is not (?:equal to )?{_value} and `x2_i` is {_value}, the result is {_value}\.$'),
     TWO_ARGS_ABSEQUAL__EQUAL = regex.compile(rf'^- +If `abs\(x1_i\)` is {_value} and `x2_i` is {_value}, the result is {_value}\.$'),
     TWO_ARGS_ABSGREATER__EQUAL = regex.compile(rf'^- +If `abs\(x1_i\)` is greater than {_value} and `x2_i` is {_value}, the result is {_value}\.$'),
     TWO_ARGS_ABSLESS__EQUAL = regex.compile(rf'^- +If `abs\(x1_i\)` is less than {_value} and `x2_i` is {_value}, the result is {_value}\.$'),