Allow scalar arguments to where()

Author: betatim

src/array_api_stubs/_draft/searching_functions.py

@@ -1,7 +1,7 @@
 __all__ = ["argmax", "argmin", "nonzero", "searchsorted", "where"]
 
 
-from ._types import Optional, Tuple, Literal, array
+from ._types import Optional, Tuple, Literal, Union, array
 
 
 def argmax(x: array, /, *, axis: Optional[int] = None, keepdims: bool = False) -> array:
@@ -139,21 +139,24 @@ def searchsorted(
     """
 
 
-def where(condition: array, x1: array, x2: array, /) -> array:
+def where(condition: array, x1: Union[array, int, float, bool], x2: Union[array, int, float, bool], /) -> array:
     """
     Returns elements chosen from ``x1`` or ``x2`` depending on ``condition``.
 
     Parameters
     ----------
     condition: array
-        when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
-    x1: array
-        first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
-    x2: array
-        second input array. Must be compatible with ``condition`` and ``x1`` (see :ref:`broadcasting`).
+        when ``True``, yield ``x1_i`` (scalar ``x1``); otherwise, yield ``x2_i`` (scalar ``x2``). Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
+    x1: Union[array, int, float, bool]
+        first input array or scalar. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
+    x2: Union[array, int, float, bool]
+        second input array or scalar. Must be compatible with ``condition`` and ``x1`` (see :ref:`broadcasting`).
 
     Returns
     -------
     out: array
         an array with elements from ``x1`` where ``condition`` is ``True``, and elements from ``x2`` elsewhere. The returned array must have a data type determined by :ref:`type-promotion` rules with the arrays ``x1`` and ``x2``.
+
+    .. versionchanged:: 2024.12
+        ``x1`` and ``x2`` may be scalars.
     """