Update dtype strictness for complex numbers in array_api elementwise functions

Original NumPy Commit: c866ef19c71b2d0269340ce984be42fd8de45e28

Author: asmeurer

array_api_strict/_elementwise_functions.py

@@ -3,9 +3,11 @@
 from ._dtypes import (
     _boolean_dtypes,
     _floating_dtypes,
+    _real_floating_dtypes,
     _complex_floating_dtypes,
     _integer_dtypes,
     _integer_or_boolean_dtypes,
+    _real_numeric_dtypes,
     _numeric_dtypes,
     _result_type,
 )
@@ -106,8 +108,8 @@ def atan2(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _floating_dtypes or x2.dtype not in _floating_dtypes:
-        raise TypeError("Only floating-point dtypes are allowed in atan2")
+    if x1.dtype not in _real_floating_dtypes or x2.dtype not in _real_floating_dtypes:
+        raise TypeError("Only real floating-point dtypes are allowed in atan2")
     # 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)
@@ -231,8 +233,8 @@ def ceil(x: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in ceil")
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in ceil")
     if x.dtype in _integer_dtypes:
         # Note: The return dtype of ceil is the same as the input
         return x
@@ -326,8 +328,8 @@ def floor(x: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in floor")
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in floor")
     if x.dtype in _integer_dtypes:
         # Note: The return dtype of floor is the same as the input
         return x
@@ -340,8 +342,8 @@ def floor_divide(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in floor_divide")
+    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 floor_divide")
     # 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)
@@ -354,8 +356,8 @@ def greater(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in greater")
+    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)
@@ -368,8 +370,8 @@ def greater_equal(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in greater_equal")
+    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)
@@ -426,8 +428,8 @@ def less(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in less")
+    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)
@@ -440,8 +442,8 @@ def less_equal(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in less_equal")
+    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)
@@ -498,8 +500,8 @@ def logaddexp(x1: Array, x2: Array) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _floating_dtypes or x2.dtype not in _floating_dtypes:
-        raise TypeError("Only floating-point dtypes are allowed in logaddexp")
+    if x1.dtype not in _real_floating_dtypes or x2.dtype not in _real_floating_dtypes:
+        raise TypeError("Only real floating-point dtypes are allowed in logaddexp")
     # 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)
@@ -639,8 +641,8 @@ def remainder(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in remainder")
+    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 remainder")
     # 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)
@@ -755,8 +757,8 @@ def trunc(x: Array, /) -> Array:
 
     See its docstring for more information.
     """
-    if x.dtype not in _numeric_dtypes:
-        raise TypeError("Only numeric dtypes are allowed in trunc")
+    if x.dtype not in _real_numeric_dtypes:
+        raise TypeError("Only real numeric dtypes are allowed in trunc")
     if x.dtype in _integer_dtypes:
         # Note: The return dtype of trunc is the same as the input
         return x

array_api_strict/tests/test_elementwise_functions.py

@@ -29,34 +29,34 @@ def test_function_types():
         "asin": "floating-point",
         "asinh": "floating-point",
         "atan": "floating-point",
-        "atan2": "floating-point",
+        "atan2": "real floating-point",
         "atanh": "floating-point",
         "bitwise_and": "integer or boolean",
         "bitwise_invert": "integer or boolean",
         "bitwise_left_shift": "integer",
         "bitwise_or": "integer or boolean",
         "bitwise_right_shift": "integer",
         "bitwise_xor": "integer or boolean",
-        "ceil": "numeric",
+        "ceil": "real numeric",
         "conj": "complex floating-point",
         "cos": "floating-point",
         "cosh": "floating-point",
         "divide": "floating-point",
         "equal": "all",
         "exp": "floating-point",
         "expm1": "floating-point",
-        "floor": "numeric",
-        "floor_divide": "numeric",
-        "greater": "numeric",
-        "greater_equal": "numeric",
+        "floor": "real numeric",
+        "floor_divide": "real numeric",
+        "greater": "real numeric",
+        "greater_equal": "real numeric",
         "imag": "complex floating-point",
         "isfinite": "numeric",
         "isinf": "numeric",
         "isnan": "numeric",
-        "less": "numeric",
-        "less_equal": "numeric",
+        "less": "real numeric",
+        "less_equal": "real numeric",
         "log": "floating-point",
-        "logaddexp": "floating-point",
+        "logaddexp": "real floating-point",
         "log10": "floating-point",
         "log1p": "floating-point",
         "log2": "floating-point",
@@ -70,7 +70,7 @@ def test_function_types():
         "positive": "numeric",
         "pow": "numeric",
         "real": "complex floating-point",
-        "remainder": "numeric",
+        "remainder": "real numeric",
         "round": "numeric",
         "sign": "numeric",
         "sin": "floating-point",
@@ -80,7 +80,7 @@ def test_function_types():
         "subtract": "numeric",
         "tan": "floating-point",
         "tanh": "floating-point",
-        "trunc": "numeric",
+        "trunc": "real numeric",
     }
 
     def _array_vals():