# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/08_trch.ipynb.

# %% auto 0
__all__ = ['pttensor', 'ptdevice', 'isdevice', 'istensor', 'istrc', 'istens', 'iscpu', 'ismps', 'iscuda', 'iscputensor',
           'ismpstensor']

# %% ../nbs/08_trch.ipynb 4
from typing import TypeGuard

# %% ../nbs/08_trch.ipynb 5
from typing import TypeGuard, Literal, Union, ForwardRef, TypeAlias, Iterable
from .core import isstr, isnone
try:
    from torch import Tensor , device    
except ImportError:
    # Define a placeholder type if torch is not installed
    Tensor = list
    device = str

pttensor: TypeAlias = ForwardRef('torch.Tensor', module='torch', is_class=True)
ptdevice: TypeAlias = ForwardRef('torch.device', module='torch', is_class=True)

# %% ../nbs/08_trch.ipynb 8
def isdevice(x) -> TypeGuard[ptdevice]:
    '''Check if `x` is a `torch.device`'''
    return not isnone(x) and ((not isstr(x) and hasattr(x, 'type')) or isinstance(x, device))

def istensor(x) -> TypeGuard[pttensor]:
    '''Check if `x` is a `torch.Tensor`
    
    See Also
    --------
    istrc : Alias for `istensor`
    istens : Alias for `istensor`
    '''
    return not isnone(x) and isinstance(x, Tensor) and hasattr(x, 'device')

def istrc(x) -> TypeGuard[ptdevice]:
    '''Check if `x` is a `torch.Tensor` with a `torch.device`.
    
    See Also
    --------
    istensor : Alias for `istrc`
    '''
    return istensor(x)

def istens(x) -> TypeGuard[pttensor]:
    '''Check if `x` is a `torch.Tensor` with a `torch.device`.
    
    See Also
    --------
    istensor : Check if `x` is a `torch.Tensor` with a `torch.device`
    istrc : Alias for `istensor`
    '''
    return istensor(x)



# %% ../nbs/08_trch.ipynb 10
def iscpu(s: str) -> TypeGuard[Literal['cpu']]:
    '''Check if `s` is `'cpu'`'''
    return s == 'cpu'

def ismps(s: str) -> TypeGuard[Literal['mps']]:
    '''Check if `s` is `'mps'`'''
    return s == 'mps'

def iscuda(s: str) -> TypeGuard[Literal['cuda']]:
    '''Check if `s` is `'cuda'`'''
    return s == 'cuda'

# %% ../nbs/08_trch.ipynb 12
def iscputensor(x) -> TypeGuard[pttensor]:
    '''Check if `x` is a `torch.Tensor` on cpu'''
    return istensor(x) and iscpu(x.device.type)

def ismpstensor(x) -> TypeGuard[pttensor]:
    '''Check if `x` is a `torch.Tensor` on mps'''
    return istensor(x) and ismps(x.device.type)

def ismpstensor(x) -> TypeGuard[pttensor]:
    '''Check if `x` is a `torch.Tensor` on cuda'''
    return istensor(x) and iscuda(x.device.type)
