Allow any combination of real dtypes in comparisons

This does not change == or != because the standard is currently unclear about
that so I'd like to see what happens there first.

Author: asmeurer

array_api_strict/_array_object.py

@@ -152,7 +152,13 @@ def __array__(self, dtype: None | np.dtype[Any] = None, copy: None | bool = None
     # spec in places where it either deviates from or is more strict than
     # NumPy behavior
 
-    def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_category: str, op: str) -> Array:
+    def _check_allowed_dtypes(
+            self,
+            other: bool | int | float | Array,
+            dtype_category: str,
+            op: str,
+            check_promotion: bool = True,
+    ) -> Array:
         """
         Helper function for operators to only allow specific input dtypes
 
@@ -176,7 +182,8 @@ def _check_allowed_dtypes(self, other: bool | int | float | Array, dtype_categor
         # This will raise TypeError for type combinations that are not allowed
         # to promote in the spec (even if the NumPy array operator would
         # promote them).
-        res_dtype = _result_type(self.dtype, other.dtype)
+        if check_promotion:
+            res_dtype = _result_type(self.dtype, other.dtype)
         if op.startswith("__i"):
             # Note: NumPy will allow in-place operators in some cases where
             # the type promoted operator does not match the left-hand side
@@ -604,7 +611,7 @@ def __ge__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __ge__.
         """
-        other = self._check_allowed_dtypes(other, "real numeric", "__ge__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__ge__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -638,7 +645,7 @@ def __gt__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __gt__.
         """
-        other = self._check_allowed_dtypes(other, "real numeric", "__gt__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__gt__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -692,7 +699,7 @@ def __le__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __le__.
         """
-        other = self._check_allowed_dtypes(other, "real numeric", "__le__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__le__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)
@@ -714,7 +721,7 @@ def __lt__(self: Array, other: Union[int, float, Array], /) -> Array:
         """
         Performs the operation __lt__.
         """
-        other = self._check_allowed_dtypes(other, "real numeric", "__lt__")
+        other = self._check_allowed_dtypes(other, "real numeric", "__lt__", check_promotion=False)
         if other is NotImplemented:
             return other
         self, other = self._normalize_two_args(self, other)

array_api_strict/_elementwise_functions.py

@@ -439,8 +439,6 @@ def greater(x1: Array, x2: Array, /) -> Array:
     """
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in greater")
-    # 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.greater(x1._array, x2._array))
 
@@ -453,8 +451,6 @@ def greater_equal(x1: Array, x2: Array, /) -> Array:
     """
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in greater_equal")
-    # 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.greater_equal(x1._array, x2._array))
 
@@ -524,8 +520,6 @@ def less(x1: Array, x2: Array, /) -> Array:
     """
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in less")
-    # 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.less(x1._array, x2._array))
 
@@ -538,8 +532,6 @@ def less_equal(x1: Array, x2: Array, /) -> Array:
     """
     if x1.dtype not in _real_numeric_dtypes or x2.dtype not in _real_numeric_dtypes:
         raise TypeError("Only real numeric dtypes are allowed in less_equal")
-    # 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.less_equal(x1._array, x2._array))