feat: add scalar support to where

PR-URL: #860 Ref: #807 Co-authored-by: Athan Reines <[email protected]> Reviewed-by: Athan Reines <[email protected]> Reviewed-by: Evgeni Burovski Reviewed-by: Lucas Colley <[email protected]>
data-apis · Jan 9, 2025 · d12a5e3 · d12a5e3
1 parent fd6f507
commit d12a5e3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 3 deletions.
diff --git a/spec/draft/API_specification/type_promotion.rst b/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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

diff --git a/src/array_api_stubs/_draft/searching_functions.py b/src/array_api_stubs/_draft/searching_functions.py
@@ -168,21 +168,35 @@ def searchsorted(
     """
 
 
-def where(condition: array, x1: array, x2: array, /) -> array:
+def where(
+    condition: array,
+    x1: Union[array, int, float, complex, bool],
+    x2: Union[array, int, float, complex, bool],
+    /,
+) -> array:
     """
     Returns elements chosen from ``x1`` or ``x2`` depending on ``condition``.
 
     Parameters
     ----------
     condition: array
         when ``True``, yield ``x1_i``; otherwise, yield ``x2_i``. Should have a boolean data type. Must be compatible with ``x1`` and ``x2`` (see :ref:`broadcasting`).
-    x1: array
+    x1: Union[array, int, float, complex, bool]
         first input array. Must be compatible with ``condition`` and ``x2`` (see :ref:`broadcasting`).
-    x2: array
+    x2: Union[array, int, float, complex, bool]
         second input array. 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
+    -----
+
+    -   At least one of  ``x1`` and ``x2`` must be an array.
+    -   If either ``x1`` or ``x2`` is a scalar value, the returned array must have a data type determined according to :ref:`mixing-scalars-and-arrays`.
+
+    .. versionchanged:: 2024.12
+        Added support for scalar arguments.
     """