Repository: data-apis/array-api-strict

Issue:
`where` not following Python scalar promotion rules?
According to the standard for [`where`](https://data-apis.org/array-api/latest/API_specification/generated/array_api.where.html), 

![Image](https://github.com/user-attachments/assets/0a43a68b-a248-4166-b466-0b8be075a279)

`array_api_strict.where` does not seem to implement this rule.

```python3
import array_api_strict as xp
x = xp.asarray([1., 2.])
xp.where(x > 1.5, x, 0)
# TypeError: array_api_strict.float64 and array_api_strict.int64 cannot be type promoted together
```

Hints from issue comments:
@ev-br: So the relevant wording from the spec is:

```
Using Python scalars (i.e., instances of bool, int, float, complex) together with arrays must be supported for:

array <op> scalar


... and scalar has a type and value compatible with the array data type:


a Python int or float for real-valued floating-point array data types.
```

so indeed  `where` is too strict here. PR welcome!


@ev-br: What I do wonder though is that `array <op> scalar` is more permissive than `array <op> array`:

```
In [9]: x = xp.asarray(1.)

In [10]: x + 0
Out[10]: Array(1., dtype=array_api_strict.float64)

In [11]: x + xp.asarray(0)
---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
...
TypeError: array_api_strict.float64 and array_api_strict.int64 cannot be type promoted together

```

This was definitely discussed and rejected (at a guess, it's granularity of floating point numbers gets larger than 1 at the end for values >> 1). Anyhow, it's covered by "A conforming implementation of the array API standard may support additional type promotion rules beyond those described in this specification."

@ev-br: > PR welcome!

Well it's easiest to just fix it, now that I very recently worked on a similar support for other binary functions, and still remember the details. 

gh-132 seems to do the trick.@mdhaber: I also ran into:
```python3
import array_api_strict as xp
xp.where(xp.asarray([True, False]), 1., xp.asarray([2, 2], dtype=xp.float32)
Array([1., 2.], dtype=array_api_strict.float64)  # should be float32
```
Will gh-132 also fix this?
Is there reason to believe the problem is isolated to `where`? I tested `clip` (very lightly) and it seems OK, and I notice that `where` was a different PR from the rest.@ev-br: > Will https://github.com/data-apis/array-api-strict/pull/132 also fix this?

Yes. It basically makes `where` reuse scalar handling from binary functions. 

```
In [1]: import array_api_strict as xp

In [2]: xp.where(xp.asarray([True, False]), 1., xp.asarray([2, 2], dtype=xp.float32))
Out[2]: Array([1., 2.], dtype=array_api_strict.float32)

In [3]: xp.__version__
Out[3]: '2.4.dev3+g33b4bf6.d20250302'
```@mdhaber: Thanks. I think I am able to remove all the test skips related to this issue now.

Solution patch:
diff --git a/array_api_strict/_helpers.py b/array_api_strict/_helpers.py
index 2258d29..d3fc9c9 100644
--- a/array_api_strict/_helpers.py
+++ b/array_api_strict/_helpers.py
@@ -31,7 +31,7 @@ def _maybe_normalize_py_scalars(x1, x2, dtype_category, func_name):
         x2 = x1._promote_scalar(x2)
     else:
         if x1.dtype not in _allowed_dtypes or x2.dtype not in _allowed_dtypes:
-            raise TypeError(f"Only {dtype_category} dtypes are allowed {func_name}. "
+            raise TypeError(f"Only {dtype_category} dtypes are allowed in {func_name}(...). "
                             f"Got {x1.dtype} and {x2.dtype}.")
     return x1, x2
 
diff --git a/array_api_strict/_searching_functions.py b/array_api_strict/_searching_functions.py
index ad32aaa..9864132 100644
--- a/array_api_strict/_searching_functions.py
+++ b/array_api_strict/_searching_functions.py
@@ -3,6 +3,7 @@ from __future__ import annotations
 from ._array_object import Array
 from ._dtypes import _result_type, _real_numeric_dtypes, bool as _bool
 from ._flags import requires_data_dependent_shapes, requires_api_version, get_array_api_strict_flags
+from ._helpers import _maybe_normalize_py_scalars
 
 from typing import TYPE_CHECKING
 if TYPE_CHECKING:
@@ -101,18 +102,7 @@ def where(
     See its docstring for more information.
     """
     if get_array_api_strict_flags()['api_version'] > '2023.12':
-        num_scalars = 0
-
-        if isinstance(x1, (bool, float, complex, int)):
-            x1 = Array._new(np.asarray(x1), device=condition.device)
-            num_scalars += 1
-
-        if isinstance(x2, (bool, float, complex, int)):
-            x2 = Array._new(np.asarray(x2), device=condition.device)
-            num_scalars += 1
-
-        if num_scalars == 2:
-            raise ValueError("One of x1, x2 arguments must be an array.")
+        x1, x2 = _maybe_normalize_py_scalars(x1, x2, "all", "where")
 
     # Call result type here just to raise on disallowed type combinations
     _result_type(x1.dtype, x2.dtype)


Test patch:
diff --git a/array_api_strict/tests/test_searching_functions.py b/array_api_strict/tests/test_searching_functions.py
index 0e54d5f..2a3a79e 100644
--- a/array_api_strict/tests/test_searching_functions.py
+++ b/array_api_strict/tests/test_searching_functions.py
@@ -20,5 +20,25 @@ def test_where_with_scalars():
     assert xp.all(x_where == expected)
 
     # The spec does not allow both x1 and x2 to be scalars
-    with pytest.raises(ValueError, match="One of"):
+    with pytest.raises(TypeError, match="Two scalars"):
         xp.where(x == 1, 42, 44)
+
+
+def test_where_mixed_dtypes():
+    # https://github.com/data-apis/array-api-strict/issues/131
+    x =  xp.asarray([1., 2.])
+    res = xp.where(x > 1.5, x, 0)
+    assert res.dtype == x.dtype
+    assert all(res == xp.asarray([0., 2.]))
+
+    # retry with boolean x1, x2
+    c = x > 1.5
+    res = xp.where(c, False, c)
+    assert all(res == xp.asarray([False, False]))
+
+
+def test_where_f32():
+    # https://github.com/data-apis/array-api-strict/issues/131#issuecomment-2723016294
+    res = xp.where(xp.asarray([True, False]), 1., xp.asarray([2, 2], dtype=xp.float32))
+    assert res.dtype == xp.float32
+