ENH: allow python scalars in binary elementwise functions

Author: ev-br

array_api_strict/_elementwise_functions.py

@@ -15,6 +15,7 @@
 from ._flags import requires_api_version
 from ._creation_functions import asarray
 from ._data_type_functions import broadcast_to, iinfo
+from ._helpers import _maybe_normalize_py_scalars
 
 from typing import Optional, Union
 
@@ -62,6 +63,8 @@ def add(x1: Array, x2: Array, /) -> Array:
 
     See its docstring for more information.
     """
+    x1, x2 = _maybe_normalize_py_scalars(x1, x2)
+
     if x1.device != x2.device:
         raise ValueError(f"Arrays from two different devices ({x1.device} and {x2.device}) can not be combined.")
 

array_api_strict/_helpers.py

@@ -0,0 +1,19 @@
+"""Private helper routines.
+"""
+import numpy as np
+
+_py_scalars = (bool, int, float, complex)
+
+def _maybe_normalize_py_scalars(x1, x2):
+    from ._array_object import Array
+
+    if isinstance(x1, _py_scalars):
+        if isinstance(x2, _py_scalars):
+            raise TypeError(f"Two scalars not allowed, {type(x1) = } and {type(x2) =}")
+        x1 = Array._new(np.asarray(x1, dtype=x2.dtype._np_dtype), device=x2.device)
+    elif isinstance(x2, _py_scalars):
+        x2 = Array._new(np.asarray(x2, dtype=x1.dtype._np_dtype), device=x1.device)
+    else:
+        # nothing to do
+        pass
+    return x1, x2

array_api_strict/tests/test_elementwise_functions.py

@@ -202,3 +202,34 @@ def test_bitwise_shift_error():
     assert_raises(
         ValueError, lambda: bitwise_right_shift(asarray([1, 1]), asarray([1, -1]))
     )
+
+
+def test_scalars():
+    # Test that binary functions accept (array, scalar) and (scalar, array) arguments
+    # and reject (scalar, scalar) arguments
+
+    def _sample_scalar(category):
+        if 'boolean' in category:
+            return True
+        elif 'floating-point' in category:
+            return 1.0
+        elif 'numeric' in category or 'integer' in category or 'all' in category:
+            return 1
+        else:
+            raise ValueError(f'Unknown {category = }')
+
+    for func_name, types in elementwise_function_input_types.items():
+        dtypes = _dtype_categories[types]
+        func = getattr(_elementwise_functions, func_name)
+        if nargs(func) == 2:
+            print(func_name, types, _sample_scalar(types))
+            scalar = _sample_scalar(types)
+            for dt in dtypes:
+                array = asarray(scalar, dtype=dt)
+                conv_scalar = asarray(scalar, dtype=array.dtype)
+                assert func(scalar, array) == func(conv_scalar, array)
+                assert func(array, scalar) == func(array, conv_scalar)
+
+                with assert_raises(TypeError):
+                    func(scalar, scalar)
+