Add scalar support to binary element-wise ops (#686)

* Add scalar support to binary element-wise ops

* Update `where` to use _promote_scalars utility

Author: tomwhite

cubed/array_api/dtypes.py

@@ -90,7 +90,9 @@
 
 
 # A Cubed-specific utility.
-def _upcast_integral_dtypes(x, dtype=None, *, allowed_dtypes=("numeric",), fname=None, device=None):
+def _upcast_integral_dtypes(
+    x, dtype=None, *, allowed_dtypes=("numeric",), fname=None, device=None
+):
     """Ensure the input dtype is allowed. If it's None, provide a good default dtype."""
     dtypes = __array_namespace_info__().default_dtypes(device=device)
 
@@ -116,3 +118,16 @@ def _upcast_integral_dtypes(x, dtype=None, *, allowed_dtypes=("numeric",), fname
             dtype = x.dtype
 
     return dtype
+
+
+def _promote_scalars(x1, x2, op):
+    """Promote at most one of x1 or x2 to an array from a Python scalar"""
+    x1_is_scalar = isinstance(x1, (int, float, complex, bool))
+    x2_is_scalar = isinstance(x2, (int, float, complex, bool))
+    if x1_is_scalar and x2_is_scalar:
+        raise TypeError(f"At least one of x1 and x2 must be an array in {op}")
+    elif x1_is_scalar:
+        x1 = x2._promote_scalar(x1)
+    elif x2_is_scalar:
+        x2 = x1._promote_scalar(x2)
+    return x1, x2

cubed/array_api/elementwise_functions.py

@@ -8,6 +8,7 @@
     _integer_dtypes,
     _integer_or_boolean_dtypes,
     _numeric_dtypes,
+    _promote_scalars,
     _real_floating_dtypes,
     _real_numeric_dtypes,
     complex64,
@@ -44,6 +45,7 @@ def acosh(x, /):
 
 
 def add(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "add")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in add")
     return elemwise(nxp.add, x1, x2, dtype=result_type(x1, x2))
@@ -68,6 +70,7 @@ def atan(x, /):
 
 
 def atan2(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "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")
     return elemwise(nxp.atan2, x1, x2, dtype=result_type(x1, x2))
@@ -80,6 +83,7 @@ def atanh(x, /):
 
 
 def bitwise_and(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_and")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -95,12 +99,14 @@ def bitwise_invert(x, /):
 
 
 def bitwise_left_shift(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_left_shift")
     if x1.dtype not in _integer_dtypes or x2.dtype not in _integer_dtypes:
         raise TypeError("Only integer dtypes are allowed in bitwise_left_shift")
     return elemwise(nxp.bitwise_left_shift, x1, x2, dtype=result_type(x1, x2))
 
 
 def bitwise_or(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_or")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -110,12 +116,14 @@ def bitwise_or(x1, x2, /):
 
 
 def bitwise_right_shift(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_right_shift")
     if x1.dtype not in _integer_dtypes or x2.dtype not in _integer_dtypes:
         raise TypeError("Only integer dtypes are allowed in bitwise_right_shift")
     return elemwise(nxp.bitwise_right_shift, x1, x2, dtype=result_type(x1, x2))
 
 
 def bitwise_xor(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "bitwise_xor")
     if (
         x1.dtype not in _integer_or_boolean_dtypes
         or x2.dtype not in _integer_or_boolean_dtypes
@@ -172,6 +180,7 @@ def conj(x, /):
 
 
 def copysign(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "copysign")
     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 copysign")
     return elemwise(nxp.copysign, x1, x2, dtype=result_type(x1, x2))
@@ -190,6 +199,7 @@ def cosh(x, /):
 
 
 def divide(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "divide")
     if x1.dtype not in _floating_dtypes or x2.dtype not in _floating_dtypes:
         raise TypeError("Only floating-point dtypes are allowed in divide")
     return elemwise(nxp.divide, x1, x2, dtype=result_type(x1, x2))
@@ -208,6 +218,7 @@ def expm1(x, /):
 
 
 def equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "equal")
     return elemwise(nxp.equal, x1, x2, dtype=nxp.bool)
 
 
@@ -221,20 +232,24 @@ def floor(x, /):
 
 
 def floor_divide(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "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")
     return elemwise(nxp.floor_divide, x1, x2, dtype=result_type(x1, x2))
 
 
 def greater(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "greater")
     return elemwise(nxp.greater, x1, x2, dtype=nxp.bool)
 
 
 def greater_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "greater_equal")
     return elemwise(nxp.greater_equal, x1, x2, dtype=nxp.bool)
 
 
 def hypot(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "hypot")
     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 hypot")
     return elemwise(nxp.hypot, x1, x2, dtype=result_type(x1, x2))
@@ -269,10 +284,12 @@ def isnan(x, /):
 
 
 def less(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "less")
     return elemwise(nxp.less, x1, x2, dtype=nxp.bool)
 
 
 def less_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "less_equal")
     return elemwise(nxp.less_equal, x1, x2, dtype=nxp.bool)
 
 
@@ -301,12 +318,14 @@ def log10(x, /):
 
 
 def logaddexp(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "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")
     return elemwise(nxp.logaddexp, x1, x2, dtype=result_type(x1, x2))
 
 
 def logical_and(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_and")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_and")
     return elemwise(nxp.logical_and, x1, x2, dtype=nxp.bool)
@@ -319,30 +338,35 @@ def logical_not(x, /):
 
 
 def logical_or(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_or")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_or")
     return elemwise(nxp.logical_or, x1, x2, dtype=nxp.bool)
 
 
 def logical_xor(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "logical_xor")
     if x1.dtype not in _boolean_dtypes or x2.dtype not in _boolean_dtypes:
         raise TypeError("Only boolean dtypes are allowed in logical_xor")
     return elemwise(nxp.logical_xor, x1, x2, dtype=nxp.bool)
 
 
 def maximum(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "maximum")
     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 maximum")
     return elemwise(nxp.maximum, x1, x2, dtype=result_type(x1, x2))
 
 
 def minimum(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "minimum")
     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 minimum")
     return elemwise(nxp.minimum, x1, x2, dtype=result_type(x1, x2))
 
 
 def multiply(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "multiply")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in multiply")
     return elemwise(nxp.multiply, x1, x2, dtype=result_type(x1, x2))
@@ -355,6 +379,7 @@ def negative(x, /):
 
 
 def not_equal(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "not_equal")
     return elemwise(nxp.not_equal, x1, x2, dtype=nxp.bool)
 
 
@@ -365,6 +390,7 @@ def positive(x, /):
 
 
 def pow(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "pow")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in pow")
     return elemwise(nxp.pow, x1, x2, dtype=result_type(x1, x2))
@@ -381,6 +407,7 @@ def real(x, /):
 
 
 def remainder(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "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")
     return elemwise(nxp.remainder, x1, x2, dtype=result_type(x1, x2))
@@ -429,6 +456,7 @@ def square(x, /):
 
 
 def subtract(x1, x2, /):
+    x1, x2 = _promote_scalars(x1, x2, "subtract")
     if x1.dtype not in _numeric_dtypes or x2.dtype not in _numeric_dtypes:
         raise TypeError("Only numeric dtypes are allowed in subtract")
     return elemwise(nxp.subtract, x1, x2, dtype=result_type(x1, x2))

cubed/array_api/searching_functions.py

@@ -1,6 +1,6 @@
 from cubed.array_api.creation_functions import asarray, zeros_like
 from cubed.array_api.data_type_functions import result_type
-from cubed.array_api.dtypes import _real_numeric_dtypes
+from cubed.array_api.dtypes import _promote_scalars, _real_numeric_dtypes
 from cubed.array_api.manipulation_functions import reshape
 from cubed.array_api.statistical_functions import max
 from cubed.backend_array_api import namespace as nxp
@@ -88,13 +88,6 @@ def _searchsorted(x, y, side):
 
 
 def where(condition, x1, x2, /):
-    x1_is_scalar = isinstance(x1, (int, float, complex, bool))
-    x2_is_scalar = isinstance(x2, (int, float, complex, bool))
-    if x1_is_scalar and x2_is_scalar:
-        raise TypeError("At least one of x1 and x2 must be an array in where")
-    elif x1_is_scalar:
-        x1 = x2._promote_scalar(x1)
-    elif x2_is_scalar:
-        x2 = x1._promote_scalar(x2)
+    x1, x2 = _promote_scalars(x1, x2, "where")
     dtype = result_type(x1, x2)
     return elemwise(nxp.where, condition, x1, x2, dtype=dtype)

cubed/tests/test_array_api.py

@@ -194,6 +194,19 @@ def test_add_different_chunks_fail(spec, executor):
     assert_array_equal(c.compute(executor=executor), np.ones((10,)) + np.ones((10,)))
 
 
+def test_add_scalars():
+    a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2))
+
+    b = xp.add(a, 1)
+    assert_array_equal(b.compute(), np.array([[2, 3, 4], [5, 6, 7], [8, 9, 10]]))
+
+    c = xp.add(2, a)
+    assert_array_equal(c.compute(), np.array([[3, 4, 5], [6, 7, 8], [9, 10, 11]]))
+
+    with pytest.raises(TypeError):
+        xp.add(1, 2)
+
+
 @pytest.mark.parametrize(
     "min, max",
     [