Skip to content

Commit

Permalink
Correct type hints in to_device and dlpack
Browse files Browse the repository at this point in the history
Signed-off-by: nstarman <nstarman@users.noreply.github.com>
  • Loading branch information
nstarman committed Nov 24, 2024
1 parent 2feaad9 commit 950fe94
Showing 1 changed file with 3 additions and 3 deletions.
6 changes: 3 additions & 3 deletions src/array_api_stubs/_draft/array_object.py
Original file line number Diff line number Diff line change
Expand Up @@ -311,7 +311,7 @@ def __dlpack__(
self,
/,
*,
stream: int | Any | None = None,
stream: Any | None = None,
max_version: tuple[int, int] | None = None,
dl_device: tuple[Enum, int] | None = None,
copy: bool | None = None,
Expand Down Expand Up @@ -1247,7 +1247,7 @@ def __xor__(self: array, other: int | bool | array, /) -> array:
...

def to_device(
self: array, device: "Device", /, *, stream: int | Any | None = None # type: ignore[type-var]
self: array, device: "Device", /, *, stream: Any | None = None # type: ignore[type-var]
) -> array:
"""
Copy the array from the device on which it currently resides to the specified ``device``.
Expand All @@ -1258,7 +1258,7 @@ def to_device(
array instance.
device: device
a ``device`` object (see :ref:`device-support`).
stream: Optional[Union[int, Any]]
stream: Optional[Any]
stream object to use during copy. In addition to the types supported in :meth:`Array.__dlpack__`, implementations may choose to support any library-specific stream object with the caveat that any code using such an object would not be portable.
Returns
Expand Down

0 comments on commit 950fe94

Please sign in to comment.