Skip to content

Commit dfea87d

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

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
@@ -294,7 +294,7 @@ def __complex__(self, /) -> complex:
294294
...
295295

296296
def __dlpack__(
297-
self, /, *, stream: int | Any | None = None
297+
self, /, *, stream: Any | None = None
298298
) -> PyCapsule: # type: ignore[type-var]
299299
"""
300300
Exports the array for consumption by :func:`~array_api.from_dlpack` as a DLPack capsule.
@@ -1076,7 +1076,7 @@ def __xor__(self: array, other: int | bool | array, /) -> array:
10761076
...
10771077

10781078
def to_device(
1079-
self: array, device: "Device", /, *, stream: int | Any | None = None # type: ignore[type-var]
1079+
self: array, device: "Device", /, *, stream: Any | None = None # type: ignore[type-var]
10801080
) -> array:
10811081
"""
10821082
Copy the array from the device on which it currently resides to the specified ``device``.
@@ -1087,7 +1087,7 @@ def to_device(
10871087
array instance.
10881088
device: device
10891089
a ``device`` object (see :ref:`device-support`).
1090-
stream: Optional[Union[int, Any]]
1090+
stream: Optional[Any]
10911091
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.
10921092
10931093
Returns

0 commit comments

Comments
 (0)