# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/06_device.ipynb.

# %% auto 0
__all__ = ['DeviceCPU', 'DeviceCUDA', 'DeviceIPU', 'DeviceXPU', 'DeviceMKLDNN', 'DeviceOpenGL', 'DeviceOpenCL', 'DeviceIDEEP',
           'DeviceHIP', 'DeviceVE', 'DeviceFPGA', 'DeviceORT', 'DeviceXLA', 'DeviceLazy', 'DeviceVulkan', 'DeviceMPS',
           'DeviceMETA', 'DeviceHPU', 'DevicePrivateUseOne', 'DeviceMeta', 'Device']

# %% ../nbs/06_device.ipynb 4
from typing import Any

# %% ../nbs/06_device.ipynb 5
from .pttype import PTMeta, PTType
from .constants import (KNOWN_DTYPES, KNOWN_DEVICES)

# %% ../nbs/06_device.ipynb 7
class DeviceMeta(PTMeta):
    """
    Metaclass for Device. It contains logic for instance checking.
    """
    KNOWN_DEVICES = KNOWN_DEVICES

    RESTRICTIONS = {
        'mps': {'Float32'}
    }

    def __getitem__(cls, item):
        if item not in cls.KNOWN_DEVICES:
            raise TypeError(f"Unknown device: {item}. Known devices are: {cls.KNOWN_DEVICES}")
        return super().__getitem__(item)

    # def __instancecheck__(cls, instance):
    #     device = cls.__args__[0]

    #     # Check device
    #     if device is not Any and instance.device.type != device:
    #         return False

    #     # Check dtype restrictions
    #     if device in cls.RESTRICTIONS and str(instance.dtype) not in cls.RESTRICTIONS[device]:
    #         return False

    #     return True

# %% ../nbs/06_device.ipynb 8
class Device(PTType, metaclass=DeviceMeta):
    """
    Represents the device a tensor is on.
    """

    __args__ = (Any, )

    # def __instancecheck__(cls, instance):
    #     device = cls.__args__[0]

    #     # Check device
    #     if device is not Any and instance.device.type != device:
    #         return False

    #     return True

# %% ../nbs/06_device.ipynb 10
DeviceCPU = Device['cpu']
DeviceCUDA = Device['cuda']
DeviceIPU = Device['ipu']
DeviceXPU = Device['xpu']
DeviceMKLDNN = Device['mkldnn']
DeviceOpenGL = Device['opengl']
DeviceOpenCL = Device['opencl']
DeviceIDEEP = Device['ideep']
DeviceHIP = Device['hip']
DeviceVE = Device['ve']
DeviceFPGA = Device['fpga']
DeviceORT = Device['ort']
DeviceXLA = Device['xla']
DeviceLazy = Device['lazy']
DeviceVulkan = Device['vulkan']
DeviceMPS = Device['mps']
DeviceMETA = Device['meta']
DeviceHPU = Device['hpu']
DevicePrivateUseOne = Device['privateuseone']
