# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/07_tens.ipynb.

# %% auto 0
__all__ = ['MemFmt', 'PresFmt', 'has_mps', 'devchain', 'try2dev', 'asdev', 'dtypechain', 'try2dtype', 'mpsdtype', 'to', 'ten2cpu',
           'ten2mps', 'ten2cuda', 'tenchain', 'try2ten', 'asten']

# %% ../nbs/07_tens.ipynb 6
from functools import wraps

# %% ../nbs/07_tens.ipynb 8
from typing import Optional, overload, TypeAlias, Union, ForwardRef

# %% ../nbs/07_tens.ipynb 9
try: from torch import memory_format as trcmemfmt, preserve_format as trcpresfmt
except ImportError: 
    class trchmemfmt(type): ...
    class trcpresfmt(trchmemfmt): ...

MemFmt = ForwardRef('torch.memory_format', is_class=True)
PresFmt = ForwardRef('torch.preserve_format', is_class=False)

# %% ../nbs/07_tens.ipynb 11
#| export


# %% ../nbs/07_tens.ipynb 13
try: import torch
except ImportError: ...

# %% ../nbs/07_tens.ipynb 15
try: import pytorch_lightning as pl
except ImportError: ...

# %% ../nbs/07_tens.ipynb 17
from atyp import IntQ, ListQ, Tensor, Device, AnyQ, DeviceQ, BoolQ, TensorQ, TorchType
from chck import isnone, isdict, notnone, istensor, istrc, isdevice, ismps, iscuda, isdtype
from calr import attrfunc, applyfns
from nlit import __DOC__, __NAME__, __MODULE__, __ANNOTATIONS__

# %% ../nbs/07_tens.ipynb 19
from .arrs import asarr, nparray, detach, clone, cpu, numpy, todense, tolist, toarray

# %% ../nbs/07_tens.ipynb 21
def has_mps() -> bool:
    '''Checks whether or not pytorch has mps availble (version) and was built with mps in mind.'''
    maybe_mps = torch.backends.mps.is_available()
    built_mps = torch.backends.mps.is_built()
    return maybe_mps and built_mps

# %% ../nbs/07_tens.ipynb 24
def devchain(x: DeviceQ, funcs, *args, **kwargs) -> Device:
    return applyfns(x, funcs, *args, check=isdevice, **kwargs)

# %% ../nbs/07_tens.ipynb 25
def try2dev(d: DeviceQ) -> Device:
    try: return d if isdevice(d) else torch.device(d)
    except: return d

# %% ../nbs/07_tens.ipynb 27
def asdev(d: DeviceQ, *args, **kwargs) -> Device:
    if isdevice(d): return d
    return devchain(d, (try2dev,), *args, **kwargs)

# %% ../nbs/07_tens.ipynb 30
def dtypechain(x: TorchType, funcs, *args, **kwargs) -> TorchType: 
    return applyfns(x, funcs, *args, check=isdtype, **kwargs)

# %% ../nbs/07_tens.ipynb 31
def try2dtype(a: Tensor, dtype: TorchType) -> Tensor:
    try: return a.to(dtype=dtype) if istensor(a) else torch.as_tensor(asarr(a), dtype=dtype)
    except: return a

# %% ../nbs/07_tens.ipynb 32
def mpsdtype(dtype: TorchType, device: DeviceQ = None, assume_on_mps: BoolQ = True) -> TorchType:
    '''Coerces `dtype` to `torch.float32` if `dtype` is `torch.float64` and `device` is `'mps'`.
    
    Parameters
    ----------    
    dtype : any
        dtype to check against
    
    device : Device, default=None
        the device of the tensor or model from which the `dtype` comes from. If provided
        will be used to detemine whether or not to make `torch.float64`, `torch.float32`
        only if the device is actually `'mps'`.

    assume_on_mps: bool, default=True
        whether or not to assume that the device of choice is `'mps'`. Setting this to
        `True` will result in `dtype` of `torch.float64` being converted to `torch.float32`
        to try and silently fix mps errors

    Notes
    -----
    - `float64` is not availble on `'mps'`, so they are coerced to `float32`
    - `assume_on_mps = True` and `has_mps() == True` are both needed as device might not 
    be provided.

    Returns
    -------
    dtype : any
        the dtype, corrected for mps if needed
    '''
    is_float64 = (dtype == torch.float64) or (dtype == 'float64') or (dtype == float)
    if notnone(device):
        if not isdevice(device): device = asdev(device)
        is_device_mps = device.type == 'mps'
        if is_device_mps:
            assume_on_mps = True

        elif device.type == 'cuda':
            assume_on_mps = False

    if has_mps() and assume_on_mps and is_float64: return torch.float32
    return dtype

# %% ../nbs/07_tens.ipynb 37
@overload
def to(
    x, other: Tensor, nonblocking: bool = False, copy: bool = False) -> Tensor: ...
@overload
def to(
    x, dtype, nonblocking: bool = False, copy: bool = False, 
    memory_format: MemFmt = trcpresfmt) -> Tensor: ...
@overload
def to(
    x, device: Device, dtype: TorchType = None, nonblocking: bool = False, copy: bool = False, 
    memory_format: MemFmt = trcpresfmt) -> Tensor: ...
@attrfunc('to', __all=True)
def to(x, *args, **kwargs) -> Tensor:
    '''Move tensor to device.'''

# %% ../nbs/07_tens.ipynb 40
@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2cpu(t: Tensor, *args, **kwargs) -> Tensor:
    '''Given a tensor, ensures that it is on cpu.'''
    return to(t, device = asdev('cpu'))

@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2mps(t: Tensor, *args, **kwargs) -> Tensor:
    '''Given a tensor, ensures that it is on mac silicon.    
    Additionally makes sure `tensor` is `torch.float32` if  `tensor.dtype` is `torch.float64`.'''
    
    if not istensor(t): return t
    dtype = mpsdtype(t.dtype, (device := asdev('mps')), assume_on_mps = False)
    return to(t, dtype = dtype, device = device)

@wraps(to, assigned = (__MODULE__, __DOC__, __ANNOTATIONS__))
def ten2cuda(t: Tensor, *args, **kwargs) -> Tensor:
    '''Given a tensor, ensures that it is on cuda.'''
    return to(t, device = asdev('cuda'))
    

# %% ../nbs/07_tens.ipynb 45
def tenchain(x: TensorQ, funcs, *args, **kwargs) -> Tensor: 
    return applyfns(x, funcs, *args, check=istensor, **kwargs)

# %% ../nbs/07_tens.ipynb 46
def try2ten(a: TensorQ, dtype: TorchType | None = None, device: DeviceQ = None) -> Tensor:
    try: return a if istensor(a) else torch.as_tensor(asarr(a), dtype=dtype, device=device)
    except: return a

# %% ../nbs/07_tens.ipynb 48
def asten(a, *args, **kwargs) -> Tensor:
    if istensor(a):
        return applyfns(a, (to, ), *args, **kwargs)

    chain = (asarr, try2move, try2ten, to, )
    a = tenchain(a, chain, *args, **kwargs)
    
    _mps, _cpu, _cuda = kwargs.get('mps', None), kwargs.get('cpu', None), kwargs.get('cuda', None)
    if _mps or ('mps' in args or asdev('mps') in args): chain += (ten2mps, )
    if _cpu or ('cpu' in args or asdev('cpu') in args): chain += (ten2cpu, )
    if _cuda or ('cuda' in args or asdev('cuda') in args): chain += (ten2cuda, )
    
    return applyfns(a, chain, *args, **kwargs)
