Skip to content

Commit 950fe94

Browse files
committed
Correct type hints in to_device and dlpack
Signed-off-by: nstarman <[email protected]>
1 parent 2feaad9 commit 950fe94

File tree

1 file changed

+3
-3
lines changed

1 file changed

+3
-3
lines changed

src/array_api_stubs/_draft/array_object.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -311,7 +311,7 @@ def __dlpack__(
311311
self,
312312
/,
313313
*,
314-
stream: int | Any | None = None,
314+
stream: Any | None = None,
315315
max_version: tuple[int, int] | None = None,
316316
dl_device: tuple[Enum, int] | None = None,
317317
copy: bool | None = None,
@@ -1247,7 +1247,7 @@ def __xor__(self: array, other: int | bool | array, /) -> array:
12471247
...
12481248

12491249
def to_device(
1250-
self: array, device: "Device", /, *, stream: int | Any | None = None # type: ignore[type-var]
1250+
self: array, device: "Device", /, *, stream: Any | None = None # type: ignore[type-var]
12511251
) -> array:
12521252
"""
12531253
Copy the array from the device on which it currently resides to the specified ``device``.
@@ -1258,7 +1258,7 @@ def to_device(
12581258
array instance.
12591259
device: device
12601260
a ``device`` object (see :ref:`device-support`).
1261-
stream: Optional[Union[int, Any]]
1261+
stream: Optional[Any]
12621262
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.
12631263
12641264
Returns

0 commit comments

Comments
 (0)