diff --git a/cubed/array_api/searching_functions.py b/cubed/array_api/searching_functions.py index 93131db8..eeec8411 100644 --- a/cubed/array_api/searching_functions.py +++ b/cubed/array_api/searching_functions.py @@ -88,5 +88,13 @@ 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) dtype = result_type(x1, x2) return elemwise(nxp.where, condition, x1, x2, dtype=dtype) diff --git a/cubed/tests/test_array_api.py b/cubed/tests/test_array_api.py index 6d2a41d1..a0b023b2 100644 --- a/cubed/tests/test_array_api.py +++ b/cubed/tests/test_array_api.py @@ -794,6 +794,22 @@ def test_searchsorted_sorter_not_implemented(): xp.searchsorted(xp.asarray([1, 0]), xp.asarray([1]), sorter=xp.asarray([1, 0])) +def test_where_scalars(): + condition = xp.asarray( + [[True, False, True], [False, True, False], [True, False, True]], chunks=(2, 2) + ) + a = xp.asarray([[1, 2, 3], [4, 5, 6], [7, 8, 9]], chunks=(2, 2)) + + b = xp.where(condition, a, 0) + assert_array_equal(b.compute(), np.array([[1, 0, 3], [0, 5, 0], [7, 0, 9]])) + + c = xp.where(condition, 0, a) + assert_array_equal(c.compute(), np.array([[0, 2, 0], [4, 0, 6], [0, 8, 0]])) + + with pytest.raises(TypeError): + xp.where(condition, 0, 1) + + # Statistical functions