Add support for scalar arguments to xp.where

Author: betatim

array_api_strict/_flags.py

@@ -23,6 +23,8 @@
     "2022.12",
     "2023.12",
 )
+# A placeholder value for the next version
+next_supported_version = "2024.12"
 
 API_VERSION = default_version = "2023.12"
 
@@ -134,7 +136,7 @@ def set_array_api_strict_flags(
     global API_VERSION, BOOLEAN_INDEXING, DATA_DEPENDENT_SHAPES, ENABLED_EXTENSIONS
 
     if api_version is not None:
-        if api_version not in supported_versions:
+        if api_version not in supported_versions + (next_supported_version,):
             raise ValueError(f"Unsupported standard version {api_version!r}")
         if api_version == "2021.12":
             warnings.warn("The 2021.12 version of the array API specification was requested but the returned namespace is actually version 2022.12", stacklevel=2)

array_api_strict/_searching_functions.py

@@ -2,7 +2,7 @@
 
 from ._array_object import Array
 from ._dtypes import _result_type, _real_numeric_dtypes
-from ._flags import requires_data_dependent_shapes, requires_api_version
+from ._flags import requires_data_dependent_shapes, requires_api_version, get_array_api_strict_flags
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -72,12 +72,19 @@ def searchsorted(
     # x1 must be 1-D, but NumPy already requires this.
     return Array._new(np.searchsorted(x1._array, x2._array, side=side, sorter=sorter), device=x1.device)
 
-def where(condition: Array, x1: Array, x2: Array, /) -> Array:
+def where(condition: Array, x1: bool | int | float | Array, x2: bool | int | float | Array, /) -> Array:
     """
     Array API compatible wrapper for :py:func:`np.where <numpy.where>`.
 
     See its docstring for more information.
     """
+    if get_array_api_strict_flags()['api_version'] > '2023.12':
+        if isinstance(x1, (bool, float, int)):
+            x1 = Array._new(np.asarray(x1), device=condition.device)
+
+        if isinstance(x2, (bool, float, int)):
+            x2 = Array._new(np.asarray(x2), device=condition.device)
+
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)
 

array_api_strict/tests/test_searching_functions.py

@@ -0,0 +1,21 @@
+import pytest
+
+import array_api_strict as xp
+
+from array_api_strict import ArrayAPIStrictFlags
+from array_api_strict._flags import next_supported_version
+
+
+def test_where_with_scalars():
+    x = xp.asarray([1, 2, 3, 1])
+
+    # Versions up to and including 2023.12 don't support scalar arguments
+    with pytest.raises(AttributeError, match="object has no attribute 'dtype'"):
+        xp.where(x == 1, 42, 44)
+
+    # Versions after 2023.12 support scalar arguments
+    with ArrayAPIStrictFlags(api_version=next_supported_version):
+        x_where = xp.where(x == 1, 42, 44)
+
+        expected = xp.asarray([42, 44, 44, 42])
+        assert xp.all(x_where == expected)