Add device kwarg support to can_cast and result_type

[kgryte]

This PR

• resolves Device-specific type promotion rules #672 by adding device keyword argument support to can_cast and result_type.

• adds guidance indicating that, by default, device capabilities should not be considered when applying type promotion rules. Currently, the specification is mum on whether device capabilities should be considered. For both functions, when device is None, only Array API type promotion rules may be applied. When device is a device object, consideration can be made as to whether particular type promotion rules are possible. E.g., if a device only supports float32, then

  can_cast(xp.float32, xp.float64, device=<device>)

  can return false. Otherwise, users have to workaround in which they apply can_cast and then use the proposed inspection APIs (Add Array API inspection utilities #689) to determine whether a device supports the promoted type. This PR makes this operation more direct.

[rgommers]

I'm not sure I like to add the device= keyword here. As discussed in gh-672, one can already make this work by passing in an array. So why not add a note instead that in case one worries about device-specific behavior, use an array on the device rather than a dtype as the first input? That avoids the need to add the keyword to these APIs in all array libraries, for most of which this will be a no-op. And it's anyway a better idea imho, because the only way to obtain a device instance is from an array, so can_cast(x, xp.float64) should be preferred over can_cast(xp.float32, xp.float32, device=x.device). In case you want this kind of logic without having an array at hand, use can_cast(xp.empty([], dtype=..., device=...), xp.float64).

[kgryte]

@rgommers Consider result_type which can take multiple operands. Which device takes precedence in the scenario where you provide arrays belonging to different devices?

For can_cast, if we require that can_cast infer the device from from, how can you resolve the cast result independent of a device? E.g., what if I want to know if I can safely cast a GPU allocated array to a CPU allocated array of greater precision? I can use astype (given the recent PR) and catch the exception, but this still seems clunky to me.

IMO, by default, these APIs should return consistent values--namely, by strictly applying type promotion rules as defined in the specification--for the sake of predictability. The device kwarg would allow an opt-in to device-specific return values.

[kgryte]

Also, adding a note here would imply that returning device-specific return values was somehow the expected behavior in prior versions of the spec. I don't think that is true given the existing language.

Determines if one data type can be cast to another data type according :ref:type-promotion rules.

...from applying the type promotion rules...

If we update the spec to infer the device from input arrays and then allow returning results only in accordance with that device, I'd consider that a breaking change to the spec.

[rgommers]

Consider result_type which can take multiple operands. Which device takes precedence in the scenario where you provide arrays belonging to different devices?

It can raise or be undefined behavior to have multiple arrays on different devices. We don't allow combining arrays that aren't on the same device in any other API either.

For can_cast, if we require that can_cast infer the device from from, how can you resolve the cast result independent of a device? E.g., what if I want to know if I can safely cast a GPU allocated array to a CPU allocated array of greater precision? I can use astype (given the recent PR) and catch the exception, but this still seems clunky to me.

You can use can_cast(x.dtype, ...) to get the device-independent behavior?

Also, adding a note here would imply that returning device-specific return values was somehow the expected behavior in prior versions of the spec. I don't think that is true given the existing language. ... I'd consider that a breaking change to the spec.

I don't quite agree. It is ill-specified now, it says nothing of relevance either way in case a dtype is missing on a specific device. We didn't consider that case at all until recently. It is and remains a bit of a corner case, so we're clarifying the expectation here. It doesn't seem reasonable to me to take a lack of precision in previous releases for a corner case to extrapolate that we must add non-useful keywords to libraries like numpy.

[kgryte]

@oleksandr-pavlyk Did you have further thoughts you wanted to share either here or on #672?

[oleksandr-pavlyk]

The type promotion rules result in a graph where dtypes are the nodes. Functions can_cast and result_type query this graph.

Some data types may be unsupported on certain devices, changing the graph. The device keyword in these functions aim to select the graph to use in the query.

What remains to clarify is how to aggregate device information from the device keyword and from input arrays (which also carry device information).

1. can_cast(from_: dtype, to_: dtype, device=None) uses full graph (as if the device supports all data types in the specification.

2. can_cast(from_: dtype, to_: dtype, device=<device>) uses promotion graph associated with the specified device

3. can_cast(from_: array, to_: dtype, device=None) uses graph corresponding to the device of input array

4. result_types(*dtypes, device=None) uses full graph

5. result_type(*arrays, device=None) requires all arrays to have the same device, or exception is raised, uses promotion rules applicable for that device.

6. result_type(*array_or_dtypes, device=<device>) requires that given device be the device where arrays were created, or exception is raised, uses promotion rules applicable for that device.

I think the spec should also provide some rules about type promotion graphs for a proper subset of dtypes.

[rgommers]

Those rules seem reasonable to me.

[kgryte]

@rgommers To be clear, you originally raised objections to adding device kwarg support. Does #691 (comment) sufficiently address your concerns? If so, I'll update this PR accordingly so that we can move things forward.

[rgommers]

Actually, re-reading the whole discussion, I think no one answered my comments around this already being possible without adding device keywords to these functions. @oleksandr-pavlyk wrote on the linked issue:

It is the use case can_cast(dtype, xp.float64), or rather result_type(xp.int64, xp.uint64) that I worry about. This should promote to the default floating point data type which is device specific.

There is no way to get a device object in a library-independent way other than taking it from an array. So this use case may not be relevant, and if it is then I'd say that it's as easy to do:

can_cast(xp.float32, xp.float64, x)

or

can_cast(xp.float32, xp.float64, asarray([], device=_my_current_device))

as it is to do

can_cast(xp.float32, xp.float64, device=_my_current_device)

Hence, unless I am missing something, this is still not compelling and already supported - we should reject this change I believe.

[kgryte]

As discussed in #672 (comment), current consensus is to not add a device kwarg to can_cast and result_type.

However, we will add, in a separate PR, some clarification to the text regarding how the respective functions should behave when provided input arrays, rather than dtypes, where the array device should be taken into consideration and reflect the device-specific type promotion graph.

Closing this PR...