# AUTOGENERATED! DO NOT EDIT! File to edit: nbs/00_torch_utils.ipynb (unless otherwise specified).

__all__ = ['isin', 'Normalizer', 'to_device', 'detach', 'is_models_equal', 'MultiWrapper', 'MultiSequential']

# Cell
import torch
import torch.nn as nn
from torch.utils.data import DataLoader

from fastcore.all import *

# Cell
def isin(t, ids):
    ''' Returns ByteTensor where True values are positions that contain ids. '''
    return (t[..., None] == torch.tensor(ids, device=t.device)).any(-1)

# Cell
class Normalizer():
    def __init__(self, device='cpu'):
        self.mean = torch.tensor([0.5, 0.5, 0.5], device=device)[None, ..., None, None] # (1, 3, 1, 1)
        self.std = torch.tensor([0.5, 0.5, 0.5], device=device)[None, ..., None, None]
    def set_device(device='cpu'):
        self.mean.to(device)
        self.std.to(device)
    def encode(self, x):
        "x: (bs, 3, _, _)"
        return (x.float()/255-self.mean) / self.std
    def decode(self, x):
        x = x*self.std + self.mean
        return (x.clamp(0., 1.)*255).long()

# Cell
def to_device(tensors, device='cpu'):
    return [t.to(device) for t in tensors]
def detach(tensors, is_to_cpu=False):
    return [t.cpu().detach() if is_to_cpu else t.detach() for t in tensors]
def is_models_equal(model_1, model_2):
    models_differ = 0
    for key_item_1, key_item_2 in zip(model_1.state_dict().items(), model_2.state_dict().items()):
        if torch.equal(key_item_1[1], key_item_2[1]):
            pass
        else:
            models_differ += 1
            if (key_item_1[0] == key_item_2[0]):
                print('Mismtach found at', key_item_1[0])
                return False
            else:
                print('Oops somethings wrong')
                return False
    if models_differ == 0:
        return True
class MultiWrapper(nn.Module):
    def __init__(self, layer, n_returns=1):
        super().__init__()
        assert n_returns>=1
        self.layer = layer
        self.n_returns = n_returns
    def forward(self, x, *others):
        if self.n_returns==1:
            return self.layer(x)
        else:
            return (self.layer(x), *others[:self.n_returns-1])
class MultiSequential(nn.Sequential):
    def forward(self, *inputs):
        for module in self._modules.values():
            if type(inputs) == tuple:
                inputs = module(*inputs)
            else:
                inputs = module(inputs)
        return inputs