|
1 | 1 | from typing import Final |
2 | | -from enum import IntEnum |
3 | 2 |
|
4 | 3 | from ._dtypes import ( |
5 | 4 | DType, float32, float64, complex64, complex128, int64, |
@@ -36,34 +35,6 @@ def __hash__(self) -> int: |
36 | 35 |
|
37 | 36 | ALL_DEVICES = (CPU_DEVICE, Device("device1"), Device("device2"), NO_FLOAT64_DEVICE) |
38 | 37 |
|
39 | | -class DLDeviceType(IntEnum): |
40 | | - kDLCPU = 1 |
41 | | - kDLCUDA = 2 |
42 | | - |
43 | | - |
44 | | -_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = { |
45 | | - CPU_DEVICE: (DLDeviceType.kDLCPU, 0), |
46 | | - Device("device1"): (DLDeviceType.kDLCUDA, 0), |
47 | | - Device("device2"): (DLDeviceType.kDLCUDA, 1), |
48 | | -} |
49 | | - |
50 | | -_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = { |
51 | | - (int(device_type), device_id): logical_device |
52 | | - for logical_device, (device_type, device_id) in _DLPACK_DEVICE_FOR.items() |
53 | | -} |
54 | | - |
55 | | - |
56 | | -def _normalize_dl_device(device_type: IntEnum | int, device_id: int) -> tuple[int, int]: |
57 | | - return (int(device_type), device_id) |
58 | | - |
59 | | - |
60 | | -def _device_from_dlpack_device( |
61 | | - device_type: IntEnum | int, device_id: int |
62 | | -) -> Device: |
63 | | - return _DLPACK_DEVICE_TO_LOGICAL.get( |
64 | | - _normalize_dl_device(device_type, device_id), CPU_DEVICE |
65 | | - ) |
66 | | - |
67 | 38 |
|
68 | 39 | def check_device(device: Device | None) -> None: |
69 | 40 | if device is not None and not isinstance(device, Device): |
|
0 commit comments