# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/05_dtype.ipynb.

# %% auto 0
__all__ = ['DTypeMeta', 'DType']

# %% ../nbs/05_dtype.ipynb 4
from typing import Any

# %% ../nbs/05_dtype.ipynb 5
import torch
from torch.optim import Optimizer
from torch.nn.modules.loss import _Loss
from torch.nn import Module

# %% ../nbs/05_dtype.ipynb 6
from .pttype import PTMeta, PTType
from .constants import (KNOWN_DTYPES, KNOWN_DEVICES)

# %% ../nbs/05_dtype.ipynb 8
class DTypeMeta(PTMeta):
    """
    Metaclass for DType. It contains logic for instance checking.
    """
    KNOWN_DTYPES = KNOWN_DTYPES
    
    def __getitem__(cls, item):
        if item not in cls.KNOWN_DTYPES:
            raise TypeError(f"Unknown dtype: {item}. Known dtypes are: {cls.KNOWN_DTYPES}")        
        return super().__getitem__(item)
    

# %% ../nbs/05_dtype.ipynb 9
class DType(PTType, metaclass=DTypeMeta):
    """
    Represents the dtype of a tensor.
    """
    __args__ = (Any, )

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

    #     # Check dtype
    #     if dtype is not Any and str(instance.dtype) != dtype:
    #         return False

    #     return True
