feat: add scalar support to where

spec/draft/API_specification/type_promotion.rst

@@ -120,6 +120,9 @@ Notes
 .. note::
    Mixed integer and floating-point type promotion rules are not specified because behavior varies between implementations.
 
+
+.. _mixing-scalars-and-arrays:
+
 Mixing arrays with Python scalars
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

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,33 @@ 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`).
+    x1: Union[array, int, float, complex, bool]
+        first input array or scalar. Scalar values are treated like an array filled with this value. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
+    x2: Union[array, int, float, complex, bool]
+        second input array or scalar. Scalar values are treated like an array filled with this value. 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``.
+
+    Notes
+    -----
+    See :ref:`mixing-scalars-and-arrays` on compatibility requirements and handling of scalar arguments for ``x1`` and ``x2``.
+
+    .. versionchanged:: 2024.12
+        ``x1`` and ``x2`` may be scalars.
     """