Closed leofang closed 2 years ago
I think the signature should be as follows:
to_device(self, /, device, *, stream=None)
where device
is a mandatory positional argument and stream
is an optional kw-only argument, with the same semantics as in the DLPack support.
Good catch, that is indeed missing.
where
device
is a mandatory positional argument
So you don't want it to be positional-only? That's possible, because it's a new method name so every library can choose the same name. However x.to_device(device='xxx')
may be a little odd with the double "device".
Hi @rgommers
So you don't want it to be positional-only?
Ah yes, I meant it to be position-only, but I typed the signature wrong. It should read
to_device(self, device, /, *, stream=None)
My apology, @rgommers, I forgot that you already had a pending PR for this (#171). Should we merge it first, and then I'll make improvements (add the stream argument, etc) in #259?
@leofang Yeah, that would probably make the most sense. Based on #171, looks like it still may still need a couple of tweaks before merging.
Thanks, Athan, I just did!
Actually, let me ask a question about stream
here: is using the DLPack-like integer the right choice here? I think libraries will either have no support for streams, or have custom objects. For DLPack there was little choice, because a custom class instance cannot work. However intra-library it could, similar to how Device
itself works, or the dtype objects. Seems like that would be more natural.
That's a good point, Ralf. Thanks for brining it up. How about we make it a union, i.e. we accept both the DLPack-like integer and library-specific stream objects?
That makes sense to me. The "integer to stream" logic must be implemented anyway to support DLPack, so is not much of an extra burden. The actual union then becomes Union[int, _Stream]
with _Stream: Any
I guess?
It would also be good to write up here in a few sentences how often stream support in to_device
is needed. I see PyTorch doesn't have it in .to
(docs)- it does have a non_blocking
keyword though, which I assume interacts with a stream context manager.
That makes sense to me. The "integer to stream" logic must be implemented anyway to support DLPack, so is not much of an extra burden. The actual union then becomes
Union[int, _Stream]
with_Stream: Any
I guess?
Technically a Union
is not the correct type here, as (with protocol arguments being contravariant) that would imply concrete implementation must, at minimum, have support for all union members. Instead, with int
- and _Stream
-support being optional, you're in a similar situation to builtin protocols such as __abs__
, wherein one of the types is a (somewhat constrained) free variable. Expressing the latter can be done with a typevar (e.g. ~T
or ~StreamType
), though arguably they don't look quite as nice as unions.
from __future__ import annotation
from typing import Union, Protocol, Any, TypeVar
_Stream = Any
_T_contra = TypeVar("_T_contra", bound=Union[None, int, _Stream], contravariant=True)
# Concrete implementations must support all union-members
class SupportsToDevice1(Protocol):
to_device(self, /, device, *, stream: int | None | _Stream = None): ...
# Concrete implementations must support one or more of the typevar bounds
class SupportsToDevice2(Protocol[_T_contra]):
to_device(self, /, device, *, stream: None | _T_contra = None): ...
It would also be good to write up here in a few sentences how often stream support in
to_device
is needed. I see PyTorch doesn't have it in.to
(docs)- it does have anon_blocking
keyword though, which I assume interacts with a stream context manager.
Yes I think so too, so PyTorch's .to(..., nonblocking, ...)
is not much different from CuPy's asnumpy(..., stream, ...)
. They're all about synchronous/asynchronous executions, which is a subject we don't want to touch in the Array API as far as I understand.
So what I have in mind is taking one step back. We can say "if stream
is specified, the copy is done on the given stream
, otherwise it's done on the library's default stream" (whatever it is, even CPU libraries have an implicit "queue/stream"). Whether or not asynchronous execution is supported is up to the library.
Expressing the latter can be done with a typevar (e.g.
~T
or~StreamType
), though arguably they don't look quite as nice as unions.
This makes perfect sense. It doesn't belong in the API standard docs though I think, that'd be pretty confusing to the reader. Maybe the right approach is to avoid using Union
and just use text to say "or any ...".
For a library implementing this, the protocol approach makes sense. Right now we don't have a good place to put these type annotations though.
We can say "if
stream
is specified, the copy is done on the givenstream
, otherwise it's done on the library's default stream" (whatever it is, even CPU libraries have an implicit "queue/stream"). Whether or not asynchronous execution is supported is up to the library.
This sounds good.
Here we mentioned there's a
to_device()
method to enable data copy/movement across devices: https://github.com/data-apis/array-api/blob/09410675703ee275e6e5746db3aded0d37bc4b96/spec/design_topics/device_support.md#L48-L49 But we haven't defined the signature and semantics yet.