Allow all dtypes in equal/not_equal/==/!= and update tests

Also update elementwise function tests to check for disallowed type
promotions, not just disallowed mixed kind types.

Author: asmeurer

array_api_strict/_array_object.py

@@ -157,6 +157,7 @@ def _check_allowed_dtypes(
             other: bool | int | float | Array,
             dtype_category: str,
             op: str,
+            *,
             check_promotion: bool = True,
     ) -> Array:
         """
@@ -577,7 +578,7 @@ def __eq__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
         """
         # Even though "all" dtypes are allowed, we still require them to be
         # promotable with each other.
-        other = self._check_allowed_dtypes(other, "all", "__eq__")
+        other = self._check_allowed_dtypes(other, "all", "__eq__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -766,7 +767,7 @@ def __ne__(self: Array, other: Union[int, float, bool, Array], /) -> Array:
         """
         Performs the operation __ne__.
         """
-        other = self._check_allowed_dtypes(other, "all", "__ne__")
+        other = self._check_allowed_dtypes(other, "all", "__ne__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)

array_api_strict/_elementwise_functions.py

@@ -375,8 +375,6 @@ def equal(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    # Call result type here just to raise on disallowed type combinations
-    _result_type(x1.dtype, x2.dtype)
     x1, x2 = Array._normalize_two_args(x1, x2)
     return Array._new(np.equal(x1._array, x2._array))
 
@@ -707,8 +705,6 @@ def not_equal(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    # Call result type here just to raise on disallowed type combinations
-    _result_type(x1.dtype, x2.dtype)
     x1, x2 = Array._normalize_two_args(x1, x2)
     return Array._new(np.not_equal(x1._array, x2._array))
 

array_api_strict/tests/test_array_object.py

@@ -94,7 +94,7 @@ def test_validate_index():
 
 def test_operators():
     # For every operator, we test that it works for the required type
-    # combinations and raises TypeError otherwise
+    # combinations and assert_raises TypeError otherwise
     binary_op_dtypes = {
         "__add__": "numeric",
         "__and__": "integer_or_boolean",
@@ -178,25 +178,25 @@ def _array_vals():
                             # See the promotion table in NEP 47 or the array
                             # API spec page on type promotion. Mixed kind
                             # promotion is not defined.
-                            if (x.dtype == uint64 and y.dtype in [int8, int16, int32, int64]
-                                or y.dtype == uint64 and x.dtype in [int8, int16, int32, int64]
-                                or x.dtype in _integer_dtypes and y.dtype not in _integer_dtypes
-                                or y.dtype in _integer_dtypes and x.dtype not in _integer_dtypes
-                                or x.dtype in _boolean_dtypes and y.dtype not in _boolean_dtypes
-                                or y.dtype in _boolean_dtypes and x.dtype not in _boolean_dtypes
-                                or x.dtype in _floating_dtypes and y.dtype not in _floating_dtypes
-                                or y.dtype in _floating_dtypes and x.dtype not in _floating_dtypes
-                                ):
-                                assert_raises(TypeError, lambda: getattr(x, _op)(y))
+                            if (op not in comparison_ops and
+                                (x.dtype == uint64 and y.dtype in [int8, int16, int32, int64]
+                                 or y.dtype == uint64 and x.dtype in [int8, int16, int32, int64]
+                                 or x.dtype in _integer_dtypes and y.dtype not in _integer_dtypes
+                                 or y.dtype in _integer_dtypes and x.dtype not in _integer_dtypes
+                                 or x.dtype in _boolean_dtypes and y.dtype not in _boolean_dtypes
+                                 or y.dtype in _boolean_dtypes and x.dtype not in _boolean_dtypes
+                                 or x.dtype in _floating_dtypes and y.dtype not in _floating_dtypes
+                                 or y.dtype in _floating_dtypes and x.dtype not in _floating_dtypes
+                                 )):
+                                assert_raises(TypeError, lambda: getattr(x, _op)(y), _op)
                             # Ensure in-place operators only promote to the same dtype as the left operand.
                             elif (
                                 _op.startswith("__i")
                                 and result_type(x.dtype, y.dtype) != x.dtype
                             ):
                                 assert_raises(TypeError, lambda: getattr(x, _op)(y), _op)
                             # Ensure only those dtypes that are required for every operator are allowed.
-                            elif (dtypes == "all" and (x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
-                                                      or x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
+                            elif (dtypes == "all"
                                 or (dtypes == "real numeric" and x.dtype in _real_numeric_dtypes and y.dtype in _real_numeric_dtypes)
                                 or (dtypes == "numeric" and x.dtype in _numeric_dtypes and y.dtype in _numeric_dtypes)
                                 or dtypes == "integer" and x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
@@ -207,7 +207,7 @@ def _array_vals():
                             ):
                                 getattr(x, _op)(y)
                             else:
-                                assert_raises(TypeError, lambda: getattr(x, _op)(y), _op)
+                                assert_raises(TypeError, lambda: getattr(x, _op)(y), (x, _op, y))
 
     unary_op_dtypes = {
         "__abs__": "numeric",

array_api_strict/tests/test_elementwise_functions.py

@@ -1,6 +1,6 @@
 from inspect import getfullargspec, getmodule
 
-from numpy.testing import assert_raises
+from .test_array_object import assert_raises
 
 from .. import asarray, _elementwise_functions
 from .._elementwise_functions import bitwise_left_shift, bitwise_right_shift
@@ -9,6 +9,11 @@
     _boolean_dtypes,
     _floating_dtypes,
     _integer_dtypes,
+    int8,
+    int16,
+    int32,
+    int64,
+    uint64,
 )
 from .._flags import set_array_api_strict_flags
 
@@ -86,6 +91,15 @@ def nargs(func):
     "trunc": "real numeric",
 }
 
+comparison_functions = [
+    'equal',
+    'greater',
+    'greater_equal',
+    'less',
+    'less_equal',
+    'not_equal',
+]
+
 def test_missing_functions():
     # Ensure the above dictionary is complete.
     import array_api_strict._elementwise_functions as mod
@@ -115,8 +129,20 @@ def _array_vals():
             func = getattr(_elementwise_functions, func_name)
             if nargs(func) == 2:
                 for y in _array_vals():
+                    # Disallow dtypes that aren't type promotable
+                    if (func_name not in comparison_functions and
+                        (x.dtype == uint64 and y.dtype in [int8, int16, int32, int64]
+                         or y.dtype == uint64 and x.dtype in [int8, int16, int32, int64]
+                         or x.dtype in _integer_dtypes and y.dtype not in _integer_dtypes
+                         or y.dtype in _integer_dtypes and x.dtype not in _integer_dtypes
+                         or x.dtype in _boolean_dtypes and y.dtype not in _boolean_dtypes
+                         or y.dtype in _boolean_dtypes and x.dtype not in _boolean_dtypes
+                         or x.dtype in _floating_dtypes and y.dtype not in _floating_dtypes
+                         or y.dtype in _floating_dtypes and x.dtype not in _floating_dtypes
+                         )):
+                        assert_raises(TypeError, lambda: func(x, y), (func_name, x, y))
                     if x.dtype not in dtypes or y.dtype not in dtypes:
-                        assert_raises(TypeError, lambda: func(x, y))
+                        assert_raises(TypeError, lambda: func(x, y), (func_name, x, y))
             else:
                 if x.dtype not in dtypes:
                     assert_raises(TypeError, lambda: func(x))