|
3 | 3 | __all__ = ["Array"] |
4 | 4 |
|
5 | 5 | from typing import TYPE_CHECKING, Protocol, TypeVar |
6 | | -from enum import Enum |
7 | 6 | from .data_types import DType |
8 | | -from ._types import Device |
| 7 | +from ._types import Device, Enum |
9 | 8 |
|
10 | 9 | if TYPE_CHECKING: |
11 | 10 | from ._types import ( |
@@ -311,7 +310,7 @@ def __dlpack__( |
311 | 310 | self, |
312 | 311 | /, |
313 | 312 | *, |
314 | | - stream: Any | None = None, |
| 313 | + stream: Any | None = None, |
315 | 314 | max_version: tuple[int, int] | None = None, |
316 | 315 | dl_device: tuple[Enum, int] | None = None, |
317 | 316 | copy: bool | None = None, |
@@ -367,10 +366,10 @@ def __dlpack__( |
367 | 366 | if it does support that), or of a different version. |
368 | 367 | This means the consumer must verify the version even when |
369 | 368 | `max_version` is passed. |
370 | | - dl_device: Optional[tuple[enum.Enum, int]] |
| 369 | + dl_device: Optional[tuple[`enum.Enum`, int]] |
371 | 370 | the DLPack device type. Default is ``None``, meaning the exported capsule |
372 | 371 | should be on the same device as ``self`` is. When specified, the format |
373 | | - must be a 2-tuple, following that of the return value of :meth:`array.__dlpack_device__`. |
| 372 | + must be a 2-tuple, following that of the return value of :meth:`Array.__dlpack_device__`. |
374 | 373 | If the device type cannot be handled by the producer, this function must |
375 | 374 | raise ``BufferError``. |
376 | 375 |
|
@@ -492,7 +491,7 @@ def __dlpack_device__(self, /) -> tuple[Enum, int]: |
492 | 491 |
|
493 | 492 | Returns |
494 | 493 | ------- |
495 | | - device: Tuple[Enum, int] |
| 494 | + device: Tuple[enum.Enum, int] |
496 | 495 | a tuple ``(device_type, device_id)`` in DLPack format. Valid device type enum members are: |
497 | 496 |
|
498 | 497 | :: |
|
0 commit comments