Skip to content

Commit dbe38b8

Browse files
committed
add dlpack_device dunder to the F32_device
1 parent 346c59a commit dbe38b8

1 file changed

Lines changed: 2 additions & 0 deletions

File tree

array_api_strict/_devices.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -43,12 +43,14 @@ def _supported_dtypes(self) -> list[DType]:
4343
class DLDeviceType(IntEnum):
4444
kDLCPU = 1
4545
kDLCUDA = 2
46+
kDLMETAL = 8
4647

4748

4849
_DLPACK_DEVICE_FOR: Final[dict[Device, tuple[DLDeviceType, int]]] = {
4950
CPU_DEVICE: (DLDeviceType.kDLCPU, 0),
5051
Device("device1"): (DLDeviceType.kDLCUDA, 0),
5152
Device("device2"): (DLDeviceType.kDLCUDA, 1),
53+
_F32_DEVICE: (DLDeviceType.kDLMETAL, 0),
5254
}
5355

5456
_DLPACK_DEVICE_TO_LOGICAL: Final[dict[tuple[int, int], Device]] = {

0 commit comments

Comments
 (0)