Construct a TVM FFI device with given device type and index. Parameters ---------- device_type: str or int The device type or name. index: int, optional The device index. Returns ------- device: tvm_ffi.Device Examples -------- Device can b
(device_type: str | int | DLDeviceType, index: int | None = None)
| 79 | |
| 80 | |
| 81 | def device(device_type: str | int | DLDeviceType, index: int | None = None) -> Device: |
| 82 | """Construct a TVM FFI device with given device type and index. |
| 83 | |
| 84 | Parameters |
| 85 | ---------- |
| 86 | device_type: str or int |
| 87 | The device type or name. |
| 88 | |
| 89 | index: int, optional |
| 90 | The device index. |
| 91 | |
| 92 | Returns |
| 93 | ------- |
| 94 | device: tvm_ffi.Device |
| 95 | |
| 96 | Examples |
| 97 | -------- |
| 98 | Device can be used to create reflection of device by |
| 99 | string representation of the device type. |
| 100 | |
| 101 | .. code-block:: python |
| 102 | |
| 103 | import tvm_ffi |
| 104 | |
| 105 | assert tvm_ffi.device("cuda:0") == tvm_ffi.device("cuda", 0) |
| 106 | assert tvm_ffi.device("cpu:0") == tvm_ffi.device("cpu", 0) |
| 107 | |
| 108 | """ |
| 109 | # must refer to core._CLASS_DEVICE so we pick up override here |
| 110 | return core._CLASS_DEVICE(device_type, index) |
| 111 | |
| 112 | |
| 113 | __all__ = ["DLDeviceType", "Device", "Tensor", "device", "from_dlpack"] |