Correct type hints in to_device and dlpack

nstarman · nstarman · commit 950fe94a775a · 2024-11-24T16:12:55.000-05:00
Signed-off-by: nstarman &lt;nstarman@users.noreply.github.com&gt;
diff --git a/src/array_api_stubs/_draft/array_object.py b/src/array_api_stubs/_draft/array_object.py
@@ -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,
@@ -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``.
@@ -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
