feat: add scalar support to where

[betatim]

Towards #807

This adds wording to the doc string and function signature to allow scalars in addition to arrays for the second and third argument.

Not super happy with the phrasing for the description of condition, maybe someone else has a suggestion for how to explain it without using a lot of words. We can then hopefully reuse that in the description of out.

There are a lot more functions that need updating (#807 (comment)). I think it make sense to get this one done, and then copy&paste for the others (instead of having a giant diff while discussing things).

There is data-apis/array-api-strict#78 which implements this in array-api-strict.

[asmeurer]

The rule for scalars should match https://data-apis.org/array-api/latest/API_specification/type_promotion.html#mixing-arrays-with-python-scalars. We should generalize that section to functions, so that we can refer to it in all the function definitions.

[asmeurer]

In other words, I think trying to word different conditions for each argument based on whether each other argument is an array or scalar is too wordy and confusing. The rule should be that scalar arguments are implicitly converted into arrays (by the rules stated in the updated version of that particular section). Then we can just talk about each argument as if it were an array. We also need to state that the behavior is undefined when all arguments are scalars. The question is really how much of that needs to be repeated in each function and how much of it we can just write once in some section and refer back to (this is not a straightforward question IMO; a lot of things in the standard are repeated for each function, since that makes it easier to read).

[kgryte]

I've opened a related PR for element-wise functions: #862

[betatim]

I think trying to word different conditions for each argument based on whether each other argument is an array or scalar is too wordy and confusing.

I agree. I like your suggestion of referring to a central explanation. In particular because I think that most people's intuition about how this will work is 95% correct and for the other 5% you need a lot of words to really explain it.

In the case of where I had assumed that the first argument always has to be an array. So the "all arguments are scalars" case can't happen.

Maybe the way to do this is a "Notes" section in the doc string that says something like "For the rules on how to handle scalar arguments see link_to_central_place." That central place could be https://data-apis.org/array-api/latest/API_specification/type_promotion.html#mixing-arrays-with-python-scalars, maybe it needs a bit of generalising to remove/rewrite the scalar OP array/array OP scalar bit.

[kgryte]

I made a few suggestions. Namely,

1. similar to feat: add scalar support to element-wise functions #862, we can keep normative language to assume array arguments. Only the type annotations need to be updated to convey scalar argument support.

2. we can allow condition to be a scalar provided that at least one of condition, x1, or x2 is an array.

3. in the case where both x1 and x2 are scalar arguments, the resulting data type should be equivalent to

   xp.result_type(xp.asarray(x1), xp.asarray(x2))

4. when only one of x1 or x2 is an array, then we can refer to the document on mixing arrays and scalars.

[rgommers]

• we can allow condition to be a scalar provided that at least one of condition, x1, or x2 is an array.

• in the case where both x1 and x2 are scalar arguments, the resulting data type should be equivalent to

We discussed these points in the community meeting yesterday, and agreed on the following:

1. condition should be an array. Rationale: (a) if it's a Python scalar, where is reduced to a simple if-statement so it doesn't seem very useful, (b) at least PyTorch doesn't support it, (c) if a use case turns up, we can add scalar support later, while the reverse isn't possible.
2. One of x1 or x2 should be an array. Rational: otherwise there is too much flexibility in dtype determination. We require the same from binary_op(x1, x2), and where is basically a binary operation from the perspective of dtype determination.

[kgryte]

LGTM.

[lucascolley]

array or scalar?

[kgryte]

We are consistent throughout the spec, atm, in describing input array arguments as arrays. Did not see a reason to deviate here.