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

# %% auto 0
__all__ = []

# %% ../nbs/patches.ipynb 2
from packaging.version import parse

import fastai
from fastai.torch_core import _rebuild_from_type

from .imports import *

# %% ../nbs/patches.ipynb 3
_torch_version = parse(torch.__version__)
_torch_20  = parse('2.0')
_torch_113 = parse('1.13')
_torch_112 = parse('1.12')

# %% ../nbs/patches.ipynb 5
# This has been upstreamed in fastai 2.7.11
if parse(fastai.__version__) < parse('2.7.11'):
    from fastai.data.core import TfmdDL

    @patch
    def to(self:TfmdDL, device):
        self.device = device
        for tfm in self.after_batch.fs:
            # Check that tfm.to is callable as TabularPandas & transforms set tfm.to as an object
            if hasattr(tfm, 'to') and callable(tfm.to):
                tfm.to(device)
            else:
                for a in L(getattr(tfm, 'parameters', None)):
                    setattr(tfm, a, getattr(tfm, a).to(device))
        return self

# %% ../nbs/patches.ipynb 7
if parse(fastai.__version__) < parse('2.7.12'):
    @patch
    def clone(self:TensorBase, *, memory_format=None):
        cls = type(self)
        return self.as_subclass(Tensor).clone(memory_format=memory_format).as_subclass(cls)

    @patch
    def new_empty(self:TensorBase, size, *, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False):
        cls = type(self)
        if _torch_version < _torch_113 and layout is None:
            layout = torch.strided
        if _torch_version < _torch_112:
            return super(TensorBase, self).new_empty(size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory, requires_grad=requires_grad)
        return self.as_subclass(Tensor).new_empty(size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory, requires_grad=requires_grad).as_subclass(cls)

    @patch
    def new_empty(self:TensorBase, *size, dtype=None, layout=None, device=None, pin_memory=False, requires_grad=False):
        cls = type(self)
        if _torch_version < _torch_113 and layout is None:
            layout = torch.strided
        if _torch_version < _torch_112:
            return super(TensorBase, self).new_empty(*size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory, requires_grad=requires_grad)
        return self.as_subclass(Tensor).new_empty(*size, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory, requires_grad=requires_grad).as_subclass(cls)

# %% ../nbs/patches.ipynb 11
if _torch_version >= _torch_20 and parse(fastai.__version__) < parse('2.7.12'):
    from fastai.callback.training import ProgressCallback

    @patch
    def __reduce_ex__(self:TensorBase, proto):
        return super(TensorBase, self).__reduce_ex__(proto)

    @patch
    def after_batch(self:ProgressCallback):
        self.pbar.update(self.iter+1)
        if hasattr(self, 'smooth_loss'):
            self.pbar.comment = f'{self.smooth_loss.item():.4f}'

# %% ../nbs/patches.ipynb 13
if parse(fastai.__version__) < parse('2.7.13'):
    from collections.abc import MutableMapping
    from fastcore.dispatch import retain_type
    from fastai.basics import defaults
    from fastai.learner import Learner

    def apply(func, x, *args, **kwargs):
        "Apply `func` recursively to `x`, passing on args"
        if is_listy(x):
            return type(x)([apply(func, o, *args, **kwargs) for o in x])
        if isinstance(x, (dict, MutableMapping)):
            return {k: apply(func, v, *args, **kwargs) for k,v in x.items()}
        res = func(x, *args, **kwargs)
        return res if x is None else retain_type(res, x)

    def to_device(b, device=None, non_blocking=False):
        "Recursively put `b` on `device`."
        if defaults.use_cuda==False:
            device='cpu'
        elif device is None:
            device=default_device()
        def _inner(o):
            if isinstance(o,Tensor):
                return o.to(device, non_blocking=non_blocking)
            return o
        return apply(_inner, b)

    @patch
    def _set_device(self:Learner, b):
        model_device = next(self.model.parameters()).device
        dls_device = getattr(self.dls, 'device', default_device())
        if model_device == dls_device:
            return to_device(b, dls_device)
        else:
            return to_device(b, model_device)
