ENH: allow python scalars in binary elementwise functions

Allow func(array, scalar) and func(scalar, array), raise on
func(scalar, scalar) if API_VERSION>=2024.12

cross-ref data-apis/array-api#807

To make sure it is all uniform,

1. Generate all binary "ufuncs" in a uniform way, with a decorator
2. Make binary "ufuncs" follow the same logic of the binary operators
3. Reuse the test loop of Array.__binop__ for binary "ufuncs"
4. (minor) in tests, reuse canonical names for dtype categories
   ("integer or boolean" vs "integer_or_boolean")

Author: ev-br

array_api_strict/_array_object.py

@@ -234,6 +234,8 @@ def _check_device(self, other):
         elif isinstance(other, Array):
             if self.device != other.device:
                 raise ValueError(f"Arrays from two different devices ({self.device} and {other.device}) can not be combined.")
+        else:
+            raise TypeError(f"Cannot combine an Array with {type(other)}.")
 
     # Helper function to match the type promotion rules in the spec
     def _promote_scalar(self, scalar):

array_api_strict/_elementwise_functions.py

@@ -0,0 +1,37 @@
+"""Private helper routines.
+"""
+
+from ._flags import get_array_api_strict_flags
+from ._dtypes import _dtype_categories
+
+_py_scalars = (bool, int, float, complex)
+
+
+def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
+
+    flags = get_array_api_strict_flags()
+    if flags["api_version"] < "2024.12":
+        # scalars will fail at the call site
+        return x1, x2
+
+    _allowed_dtypes = _dtype_categories[dtype_category]
+
+    if isinstance(x1, _py_scalars):
+        if isinstance(x2, _py_scalars):
+            raise TypeError(f"Two scalars not allowed, got {type(x1) = } and {type(x2) =}")
+        # x2 must be an array
+        if x2.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x2.dtype}.")
+        x1 = x2._promote_scalar(x1)
+
+    elif isinstance(x2, _py_scalars):
+        # x1 must be an array
+        if x1.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. Got {x1.dtype}.")
+        x2 = x1._promote_scalar(x2)
+    else:
+        if x1.dtype not in _allowed_dtypes or x2.dtype not in _allowed_dtypes:
+            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. "
+                            f"Got {x1.dtype} and {x2.dtype}.")
+    return x1, x2
+

array_api_strict/_helpers.py

@@ -96,12 +96,60 @@ def test_promoted_scalar_inherits_device():
 
     assert y.device == device1
 
+
+BIG_INT = int(1e30)
+
+def _check_op_array_scalar(dtypes, a, s, func, func_name, BIG_INT=BIG_INT):
+    # Test array op scalar. From the spec, the following combinations
+    # are supported:
+
+    # - Python bool for a bool array dtype,
+    # - a Python int within the bounds of the given dtype for integer array dtypes,
+    # - a Python int or float for real floating-point array dtypes
+    # - a Python int, float, or complex for complex floating-point array dtypes
+
+    if ((dtypes == "all"
+         or dtypes == "numeric" and a.dtype in _numeric_dtypes
+         or dtypes == "real numeric" and a.dtype in _real_numeric_dtypes
+         or dtypes == "integer" and a.dtype in _integer_dtypes
+         or dtypes == "integer or boolean" and a.dtype in _integer_or_boolean_dtypes
+         or dtypes == "boolean" and a.dtype in _boolean_dtypes
+         or dtypes == "floating-point" and a.dtype in _floating_dtypes
+         or dtypes == "real floating-point" and a.dtype in _real_floating_dtypes
+        )
+        # bool is a subtype of int, which is why we avoid
+        # isinstance here.
+        and (a.dtype in _boolean_dtypes and type(s) == bool
+             or a.dtype in _integer_dtypes and type(s) == int
+             or a.dtype in _real_floating_dtypes and type(s) in [float, int]
+             or a.dtype in _complex_floating_dtypes and type(s) in [complex, float, int]
+        )):
+        if a.dtype in _integer_dtypes and s == BIG_INT:
+            with assert_raises(OverflowError):
+                func(s)
+            return False
+
+        else:
+            # Only test for no error
+            with suppress_warnings() as sup:
+                # ignore warnings from pow(BIG_INT)
+                sup.filter(RuntimeWarning,
+                           "invalid value encountered in power")
+                func(s)
+            return True
+
+    else:
+        with assert_raises(TypeError):
+            func(s)
+        return False
+
+
 def test_operators():
     # For every operator, we test that it works for the required type
     # combinations and raises TypeError otherwise
     binary_op_dtypes = {
         "__add__": "numeric",
-        "__and__": "integer_or_boolean",
+        "__and__": "integer or boolean",
         "__eq__": "all",
         "__floordiv__": "real numeric",
         "__ge__": "real numeric",
@@ -112,12 +160,12 @@ def test_operators():
         "__mod__": "real numeric",
         "__mul__": "numeric",
         "__ne__": "all",
-        "__or__": "integer_or_boolean",
+        "__or__": "integer or boolean",
         "__pow__": "numeric",
         "__rshift__": "integer",
         "__sub__": "numeric",
-        "__truediv__": "floating",
-        "__xor__": "integer_or_boolean",
+        "__truediv__": "floating-point",
+        "__xor__": "integer or boolean",
     }
     # Recompute each time because of in-place ops
     def _array_vals():
@@ -128,8 +176,6 @@ def _array_vals():
         for d in _floating_dtypes:
             yield asarray(1.0, dtype=d)
 
-
-    BIG_INT = int(1e30)
     for op, dtypes in binary_op_dtypes.items():
         ops = [op]
         if op not in ["__eq__", "__ne__", "__le__", "__ge__", "__lt__", "__gt__"]:
@@ -139,40 +185,7 @@ def _array_vals():
         for s in [1, 1.0, 1j, BIG_INT, False]:
             for _op in ops:
                 for a in _array_vals():
-                    # Test array op scalar. From the spec, the following combinations
-                    # are supported:
-
-                    # - Python bool for a bool array dtype,
-                    # - a Python int within the bounds of the given dtype for integer array dtypes,
-                    # - a Python int or float for real floating-point array dtypes
-                    # - a Python int, float, or complex for complex floating-point array dtypes
-
-                    if ((dtypes == "all"
-                         or dtypes == "numeric" and a.dtype in _numeric_dtypes
-                         or dtypes == "real numeric" and a.dtype in _real_numeric_dtypes
-                         or dtypes == "integer" and a.dtype in _integer_dtypes
-                         or dtypes == "integer_or_boolean" and a.dtype in _integer_or_boolean_dtypes
-                         or dtypes == "boolean" and a.dtype in _boolean_dtypes
-                         or dtypes == "floating" and a.dtype in _floating_dtypes
-                        )
-                        # bool is a subtype of int, which is why we avoid
-                        # isinstance here.
-                        and (a.dtype in _boolean_dtypes and type(s) == bool
-                             or a.dtype in _integer_dtypes and type(s) == int
-                             or a.dtype in _real_floating_dtypes and type(s) in [float, int]
-                             or a.dtype in _complex_floating_dtypes and type(s) in [complex, float, int]
-                        )):
-                        if a.dtype in _integer_dtypes and s == BIG_INT:
-                            assert_raises(OverflowError, lambda: getattr(a, _op)(s))
-                        else:
-                            # Only test for no error
-                            with suppress_warnings() as sup:
-                                # ignore warnings from pow(BIG_INT)
-                                sup.filter(RuntimeWarning,
-                                           "invalid value encountered in power")
-                                getattr(a, _op)(s)
-                    else:
-                        assert_raises(TypeError, lambda: getattr(a, _op)(s))
+                    _check_op_array_scalar(dtypes, a, s, getattr(a, _op), _op)
 
                 # Test array op array.
                 for _op in ops:
@@ -203,18 +216,18 @@ def _array_vals():
                                 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
-                                or dtypes == "integer_or_boolean" and (x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
+                                or dtypes == "integer or boolean" and (x.dtype in _integer_dtypes and y.dtype in _integer_dtypes
                                                                        or x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes)
                                 or dtypes == "boolean" and x.dtype in _boolean_dtypes and y.dtype in _boolean_dtypes
-                                or dtypes == "floating" and x.dtype in _floating_dtypes and y.dtype in _floating_dtypes
+                                or dtypes == "floating-point" and x.dtype in _floating_dtypes and y.dtype in _floating_dtypes
                             ):
                                 getattr(x, _op)(y)
                             else:
                                 assert_raises(TypeError, lambda: getattr(x, _op)(y))
 
     unary_op_dtypes = {
         "__abs__": "numeric",
-        "__invert__": "integer_or_boolean",
+        "__invert__": "integer or boolean",
         "__neg__": "numeric",
         "__pos__": "numeric",
     }
@@ -223,7 +236,7 @@ def _array_vals():
             if (
                 dtypes == "numeric"
                 and a.dtype in _numeric_dtypes
-                or dtypes == "integer_or_boolean"
+                or dtypes == "integer or boolean"
                 and a.dtype in _integer_or_boolean_dtypes
             ):
                 # Only test for no error

array_api_strict/tests/test_array_object.py

@@ -1,6 +1,7 @@
 from inspect import signature, getmodule
 
-from numpy.testing import assert_raises
+from pytest import raises as assert_raises
+from numpy.testing import suppress_warnings
 
 import pytest
 
@@ -19,6 +20,8 @@
 )
 from .._flags import set_array_api_strict_flags
 
+from .test_array_object import _check_op_array_scalar, BIG_INT
+
 import array_api_strict
 
 
@@ -120,6 +123,7 @@ def test_missing_functions():
     # Ensure the above dictionary is complete.
     import array_api_strict._elementwise_functions as mod
     mod_funcs = [n for n in dir(mod) if getmodule(getattr(mod, n)) is mod]
+    mod_funcs = [n for n in mod_funcs if not n.startswith("_")]
     assert set(mod_funcs) == set(elementwise_function_input_types)
 
 
@@ -202,3 +206,51 @@ def test_bitwise_shift_error():
     assert_raises(
         ValueError, lambda: bitwise_right_shift(asarray([1, 1]), asarray([1, -1]))
     )
+
+
+
+def test_scalars():
+    # mirror test_array_object.py::test_operators()
+    #
+    # Also check that binary functions accept (array, scalar) and (scalar, array)
+    # arguments, and reject (scalar, scalar) arguments.
+
+    # Use the latest version of the standard so that scalars are actually allowed
+    with pytest.warns(UserWarning):
+        set_array_api_strict_flags(api_version="2024.12")
+
+    def _array_vals():
+        for d in _integer_dtypes:
+            yield asarray(1, dtype=d)
+        for d in _boolean_dtypes:
+            yield asarray(False, dtype=d)
+        for d in _floating_dtypes:
+            yield asarray(1.0, dtype=d)
+
+
+    for func_name, dtypes in elementwise_function_input_types.items():
+        func = getattr(_elementwise_functions, func_name)
+        if nargs(func) != 2:
+            continue
+
+        for s in [1, 1.0, 1j, BIG_INT, False]:
+            for a in _array_vals():
+                for func1 in [lambda s: func(a, s), lambda s: func(s, a)]:
+                    allowed = _check_op_array_scalar(dtypes, a, s, func1, func_name)
+
+                    # only check `func(array, scalar) == `func(array, array)` if
+                    # the former is legal under the promotion rules
+                    if allowed:
+                        conv_scalar = asarray(s, dtype=a.dtype)
+
+                        with suppress_warnings() as sup:
+                            # ignore warnings from pow(BIG_INT)
+                            sup.filter(RuntimeWarning,
+                                       "invalid value encountered in power")
+                            assert func(s, a) == func(conv_scalar, a)
+                            assert func(a, s) == func(a, conv_scalar)
+
+                        with assert_raises(TypeError):
+                            func(s, s)
+
+