import sys
_module = sys.modules[__name__]
del sys
conf = _module
custom_nn = _module
aae = _module
betatcvae = _module
ciwae = _module
hvae = _module
iwae = _module
miwae = _module
piwae = _module
pvae = _module
rae_gp = _module
rae_l2 = _module
svae = _module
vae = _module
vamp = _module
wae = _module
training = _module
setup = _module
pythae = _module
config = _module
customexception = _module
data = _module
datasets = _module
preprocessors = _module
models = _module
adversarial_ae = _module
adversarial_ae_config = _module
adversarial_ae_model = _module
ae = _module
ae_config = _module
ae_model = _module
auto_model = _module
auto_config = _module
auto_model = _module
base = _module
base_config = _module
base_model = _module
base_utils = _module
beta_tc_vae = _module
beta_tc_vae_config = _module
beta_tc_vae_model = _module
beta_vae = _module
beta_vae_config = _module
beta_vae_model = _module
ciwae_config = _module
ciwae_model = _module
disentangled_beta_vae = _module
disentangled_beta_vae_config = _module
disentangled_beta_vae_model = _module
factor_vae = _module
factor_vae_config = _module
factor_vae_model = _module
factor_vae_utils = _module
hvae_config = _module
hvae_model = _module
info_vae = _module
info_vae_config = _module
info_vae_model = _module
iwae_config = _module
iwae_model = _module
miwae_config = _module
miwae_model = _module
msssim_vae = _module
msssim_vae_config = _module
msssim_vae_model = _module
msssim_vae_utils = _module
nn = _module
base_architectures = _module
benchmarks = _module
celeba = _module
convnets = _module
resnets = _module
cifar = _module
convnets = _module
resnets = _module
mnist = _module
convnets = _module
resnets = _module
utils = _module
default_architectures = _module
normalizing_flows = _module
base_nf_config = _module
base_nf_model = _module
iaf = _module
iaf_config = _module
iaf_model = _module
layers = _module
made = _module
made_config = _module
made_model = _module
maf = _module
maf_config = _module
maf_model = _module
pixelcnn = _module
pixelcnn_config = _module
pixelcnn_model = _module
utils = _module
planar_flow = _module
planar_flow_config = _module
planar_flow_model = _module
radial_flow = _module
radial_flow_config = _module
radial_flow_model = _module
piwae_config = _module
piwae_model = _module
pvae_config = _module
pvae_model = _module
pvae_utils = _module
rae_gp_config = _module
rae_gp_model = _module
rae_l2_config = _module
rae_l2_model = _module
rhvae = _module
rhvae_config = _module
rhvae_model = _module
rhvae_utils = _module
svae_config = _module
svae_model = _module
svae_utils = _module
vae_config = _module
vae_model = _module
vae_gan = _module
vae_gan_config = _module
vae_gan_model = _module
vae_iaf = _module
vae_iaf_config = _module
vae_iaf_model = _module
vae_lin_nf = _module
vae_lin_nf_config = _module
vae_lin_nf_model = _module
vamp_config = _module
vamp_model = _module
vq_vae = _module
vq_vae_config = _module
vq_vae_model = _module
vq_vae_utils = _module
wae_mmd = _module
wae_mmd_config = _module
wae_mmd_model = _module
pipelines = _module
base_pipeline = _module
generation = _module
pipeline_utils = _module
training = _module
samplers = _module
base_sampler = _module
base_sampler_config = _module
gaussian_mixture = _module
gaussian_mixture_config = _module
gaussian_mixture_sampler = _module
hypersphere_uniform_sampler = _module
hypersphere_uniform_config = _module
hypersphere_uniform_sampler = _module
iaf_sampler = _module
iaf_sampler = _module
iaf_sampler_config = _module
maf_sampler = _module
maf_sampler = _module
maf_sampler_config = _module
manifold_sampler = _module
rhvae_sampler = _module
rhvae_sampler_config = _module
normal_sampling = _module
normal_config = _module
normal_sampler = _module
pixelcnn_sampler = _module
pixelcnn_sampler = _module
pixelcnn_sampler_config = _module
pvae_sampler = _module
pvae_sampler = _module
pvae_sampler_config = _module
two_stage_vae_sampler = _module
two_stage_sampler = _module
two_stage_sampler_config = _module
vamp_sampler = _module
vamp_sampler = _module
vamp_sampler_config = _module
trainers = _module
adversarial_trainer = _module
adversarial_trainer = _module
adversarial_trainer_config = _module
base_trainer = _module
base_trainer = _module
base_training_config = _module
coupled_optimizer_adversarial_trainer = _module
coupled_optimizer_adversarial_trainer = _module
coupled_optimizer_adversarial_trainer_config = _module
coupled_optimizer_trainer = _module
coupled_optimizer_trainer = _module
coupled_optimizer_trainer_config = _module
trainer_utils = _module
training_callbacks = _module
tests = _module
conftest = _module
custom_architectures = _module
test_AE = _module
test_Adversarial_AE = _module
test_BetaTCVAE = _module
test_BetaVAE = _module
test_CIWAE = _module
test_DisentangledBetaVAE = _module
test_FactorVAE = _module
test_HVAE = _module
test_IAF = _module
test_IWAE = _module
test_MADE = _module
test_MAF = _module
test_MIWAE = _module
test_MSSSIMVAE = _module
test_PIWAE = _module
test_PixelCNN = _module
test_PoincareVAE = _module
test_RHVAE = _module
test_SVAE = _module
test_VAE = _module
test_VAEGAN = _module
test_VAE_IAF = _module
test_VAE_LinFlow = _module
test_VAMP = _module
test_VQVAE = _module
test_WAE_MMD = _module
test_adversarial_trainer = _module
test_auto_model = _module
test_baseAE = _module
test_baseSampler = _module
test_base_trainer = _module
test_config = _module
test_coupled_optimizers_adversarial_trainer = _module
test_coupled_optimizers_trainer = _module
test_datasets = _module
test_gaussian_mixture_sampler = _module
test_hypersphere_uniform_sampler = _module
test_iaf_sampler = _module
test_info_vae_mmd = _module
test_maf_sampler = _module
test_nn_benchmark = _module
test_normal_sampler = _module
test_pipeline_standalone = _module
test_pixelcnn_sampler = _module
test_planar_flow = _module
test_preprocessing = _module
test_pvae_sampler = _module
test_radial_flow = _module
test_rae_gp = _module
test_rae_l2 = _module
test_rhvae_sampler = _module
test_training_callbacks = _module
test_two_stage_sampler = _module
test_vamp_sampler = _module

from _paritybench_helpers import _mock_config, patch_functional
from unittest.mock import mock_open, MagicMock
from torch.autograd import Function
from torch.nn import Module
import abc, collections, copy, enum, functools, inspect, itertools, logging, math, matplotlib, numbers, numpy, pandas, queue, random, re, scipy, sklearn, string, tensorflow, time, torch, torchaudio, torchtext, torchvision, types, typing, uuid, warnings
import numpy as np
from torch import Tensor
patch_functional()
open = mock_open()
yaml = logging = sys = argparse = MagicMock()
ArgumentParser = argparse.ArgumentParser
_global_config = args = argv = cfg = config = params = _mock_config()
argparse.ArgumentParser.return_value.parse_args.return_value = _global_config
yaml.load.return_value = _global_config
sys.argv = _global_config
__version__ = '1.0.0'
xrange = range
wraps = functools.wraps


import torch


import torch.nn as nn


import numpy as np


from typing import List


import logging


from torch.utils.data import Dataset


import math


import torch.nn.functional as F


from time import time


from collections import OrderedDict


from typing import Any


from typing import Tuple


from typing import Optional


from typing import Union


import inspect


import warnings


from copy import deepcopy


from torch.autograd import grad


from numbers import Number


import torch.distributions as dist


from torch.autograd import Function


from torch.distributions import Normal


from torch.distributions.utils import _standard_normal


from torch.distributions.utils import broadcast_all


from torch.nn import functional as F


from collections import deque


import scipy.special


from sklearn import mixture


from torch.utils.data import DataLoader


from torch.distributions import MultivariateNormal


import itertools


import torch.optim as optim


from torch.optim.lr_scheduler import ReduceLROnPlateau


from typing import Dict


import random


import typing


from torch.optim import Adam


from torch.optim import SGD


from torch.optim import Adadelta


from torch.optim import Adagrad


from torch.optim import RMSprop


from sklearn import manifold


from torch.optim.lr_scheduler import StepLR


from torch.optim.lr_scheduler import LinearLR


from torch.optim.lr_scheduler import ExponentialLR


from collections import Counter


class RiemannianLayer(nn.Module):

    def __init__(self, in_features, out_features, manifold, over_param, weight_norm):
        super(RiemannianLayer, self).__init__()
        self.in_features = in_features
        self.out_features = out_features
        self.manifold = manifold
        self._weight = nn.Parameter(torch.Tensor(out_features, in_features))
        self.over_param = over_param
        self.weight_norm = weight_norm
        self._bias = nn.Parameter(torch.Tensor(out_features, 1))
        self.reset_parameters()

    @property
    def weight(self):
        return self.manifold.transp0(self.bias, self._weight)

    @property
    def bias(self):
        if self.over_param:
            return self._bias
        else:
            return self.manifold.expmap0(self._weight * self._bias)

    def reset_parameters(self):
        nn.init.kaiming_normal_(self._weight, a=math.sqrt(5))
        fan_in, _ = nn.init._calculate_fan_in_and_fan_out(self._weight)
        bound = 4 / math.sqrt(fan_in)
        nn.init.uniform_(self._bias, -bound, bound)
        if self.over_param:
            with torch.no_grad():
                self._bias.set_(self.manifold.expmap0(self._bias))


class GeodesicLayer(RiemannianLayer):

    def __init__(self, in_features, out_features, manifold, over_param=False, weight_norm=False):
        super(GeodesicLayer, self).__init__(in_features, out_features, manifold, over_param, weight_norm)

    def forward(self, input):
        input = input.unsqueeze(0)
        input = input.unsqueeze(-2).expand(*input.shape[:-(len(input.shape) - 2)], self.out_features, self.in_features)
        res = self.manifold.normdist2plane(input, self.bias, self.weight, signed=True, norm=self.weight_norm)
        return res


class NonLinear(nn.Module):

    def __init__(self, input_size, output_size, bias=True, activation=None):
        super(NonLinear, self).__init__()
        self.activation = activation
        self.linear = nn.Linear(int(input_size), int(output_size), bias=bias)

    def forward(self, x):
        h = self.linear(x)
        if self.activation is not None:
            h = self.activation(h)
        return h


class GatedDense(nn.Module):

    def __init__(self, input_size, output_size, activation=None):
        super(GatedDense, self).__init__()
        self.activation = activation
        self.sigmoid = nn.Sigmoid()
        self.h = nn.Linear(input_size, output_size)
        self.g = nn.Linear(input_size, output_size)

    def forward(self, x):
        h = self.h(x)
        if self.activation is not None:
            h = self.activation(self.h(x))
        g = self.sigmoid(self.g(x))
        return h * g


def hf_hub_is_available():
    return importlib.util.find_spec('huggingface_hub') is not None


logger = logging.getLogger(__name__)


class AutoModel(nn.Module):
    """Utils class allowing to reload any :class:`pythae.models` automatically"""

    def __init__(self) ->None:
        super().__init__()

    @classmethod
    def load_from_folder(cls, dir_path: str):
        """Class method to be used to load the model from a specific folder

        Args:
            dir_path (str): The path where the model should have been be saved.

        .. note::
            This function requires the folder to contain:

            - | a ``model_config.json`` and a ``model.pt`` if no custom architectures were provided

            **or**

            - | a ``model_config.json``, a ``model.pt`` and a ``encoder.pkl`` (resp.
                ``decoder.pkl``) if a custom encoder (resp. decoder) was provided
        """
        with open(os.path.join(dir_path, 'model_config.json')) as f:
            model_name = json.load(f)['name']
        if model_name == 'Adversarial_AE_Config':
            model = Adversarial_AE.load_from_folder(dir_path=dir_path)
        elif model_name == 'AEConfig':
            model = AE.load_from_folder(dir_path=dir_path)
        elif model_name == 'BetaTCVAEConfig':
            model = BetaTCVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'BetaVAEConfig':
            model = BetaVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'DisentangledBetaVAEConfig':
            model = DisentangledBetaVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'FactorVAEConfig':
            model = FactorVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'HVAEConfig':
            model = HVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'INFOVAE_MMD_Config':
            model = INFOVAE_MMD.load_from_folder(dir_path=dir_path)
        elif model_name == 'IWAEConfig':
            model = IWAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'MSSSIM_VAEConfig':
            model = MSSSIM_VAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'RAE_GP_Config':
            model = RAE_GP.load_from_folder(dir_path=dir_path)
        elif model_name == 'RAE_L2_Config':
            model = RAE_L2.load_from_folder(dir_path=dir_path)
        elif model_name == 'RHVAEConfig':
            model = RHVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'SVAEConfig':
            model = SVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'VAEConfig':
            model = VAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'VAEGANConfig':
            model = VAEGAN.load_from_folder(dir_path=dir_path)
        elif model_name == 'VAE_IAF_Config':
            model = VAE_IAF.load_from_folder(dir_path=dir_path)
        elif model_name == 'VAE_LinNF_Config':
            model = VAE_LinNF.load_from_folder(dir_path=dir_path)
        elif model_name == 'VAMPConfig':
            model = VAMP.load_from_folder(dir_path=dir_path)
        elif model_name == 'VQVAEConfig':
            model = VQVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'WAE_MMD_Config':
            model = WAE_MMD.load_from_folder(dir_path=dir_path)
        elif model_name == 'MAFConfig':
            model = MAF.load_from_folder(dir_path=dir_path)
        elif model_name == 'IAFConfig':
            model = IAF.load_from_folder(dir_path=dir_path)
        elif model_name == 'PlanarFlowConfig':
            model = PlanarFlow.load_from_folder(dir_path=dir_path)
        elif model_name == 'RadialFlowConfig':
            model = RadialFlow.load_from_folder(dir_path=dir_path)
        elif model_name == 'MADEConfig':
            model = MADE.load_from_folder(dir_path=dir_path)
        elif model_name == 'PixelCNNConfig':
            model = PixelCNN.load_from_folder(dir_path=dir_path)
        elif model_name == 'PoincareVAEConfig':
            model = PoincareVAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'CIWAEConfig':
            model = CIWAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'MIWAEConfig':
            model = MIWAE.load_from_folder(dir_path=dir_path)
        elif model_name == 'PIWAEConfig':
            model = PIWAE.load_from_folder(dir_path=dir_path)
        else:
            raise NameError(f'Cannot reload automatically the model... The model name in the `model_config.json may be corrupted. Got {model_name}')
        return model

    @classmethod
    def load_from_hf_hub(cls, hf_hub_path: str, allow_pickle: bool=False):
        """Class method to be used to load a automaticaly a pretrained model from the Hugging Face
        hub

        Args:
            hf_hub_path (str): The path where the model should have been be saved on the
                hugginface hub.

        .. note::
            This function requires the folder to contain:

            - | a ``model_config.json`` and a ``model.pt`` if no custom architectures were provided

            **or**

            - | a ``model_config.json``, a ``model.pt`` and a ``encoder.pkl`` (resp.
                ``decoder.pkl``) if a custom encoder (resp. decoder) was provided
        """
        if not hf_hub_is_available():
            raise ModuleNotFoundError('`huggingface_hub` package must be installed to load models from the HF hub. Run `python -m pip install huggingface_hub` and log in to your account with `huggingface-cli login`.')
        logger.info(f'Downloading config file ...')
        config_path = hf_hub_download(repo_id=hf_hub_path, filename='model_config.json')
        dir_path = os.path.dirname(config_path)
        with open(os.path.join(dir_path, 'model_config.json')) as f:
            model_name = json.load(f)['name']
        if model_name == 'Adversarial_AE_Config':
            model = Adversarial_AE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'AEConfig':
            model = AE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'BetaTCVAEConfig':
            model = BetaTCVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'BetaVAEConfig':
            model = BetaVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'DisentangledBetaVAEConfig':
            model = DisentangledBetaVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'FactorVAEConfig':
            model = FactorVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'HVAEConfig':
            model = HVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'INFOVAE_MMD_Config':
            model = INFOVAE_MMD.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'IWAEConfig':
            model = IWAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'MSSSIM_VAEConfig':
            model = MSSSIM_VAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'RAE_GP_Config':
            model = RAE_GP.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'RAE_L2_Config':
            model = RAE_L2.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'RHVAEConfig':
            model = RHVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'SVAEConfig':
            model = SVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VAEConfig':
            model = VAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VAEGANConfig':
            model = VAEGAN.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VAE_IAF_Config':
            model = VAE_IAF.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VAE_LinNF_Config':
            model = VAE_LinNF.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VAMPConfig':
            model = VAMP.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'VQVAEConfig':
            model = VQVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'WAE_MMD_Config':
            model = WAE_MMD.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'MAFConfig':
            model = MAF.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'IAFConfig':
            model = IAF.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'PlanarFlowConfig':
            model = PlanarFlow.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'RadialFlowConfig':
            model = RadialFlow.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'MADEConfig':
            model = MADE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'PixelCNNConfig':
            model = PixelCNN.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'PoincareVAEConfig':
            model = PoincareVAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'CIWAEConfig':
            model = CIWAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'MIWAEConfig':
            model = MIWAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        elif model_name == 'PIWAEConfig':
            model = PIWAE.load_from_hf_hub(hf_hub_path=hf_hub_path, allow_pickle=allow_pickle)
        else:
            raise NameError(f'Cannot reload automatically the model... The model name in the `model_config.json may be corrupted. Got {model_name}')
        return model


class BadInheritanceError(Exception):
    pass


class DatasetOutput(OrderedDict):
    """Base DatasetOutput class fixing the output type from the dataset. This class is inspired from
    the ``ModelOutput`` class from hugginface transformers library"""

    def __getitem__(self, k):
        if isinstance(k, str):
            self_dict = {k: v for k, v in self.items()}
            return self_dict[k]
        else:
            return self.to_tuple()[k]

    def __setattr__(self, name, value):
        super().__setitem__(name, value)
        super().__setattr__(name, value)

    def __setitem__(self, key, value):
        super().__setitem__(key, value)
        super().__setattr__(key, value)

    def to_tuple(self) ->Tuple[Any]:
        """
        Convert self to a tuple containing all the attributes/keys that are not ``None``.
        """
        return tuple(self[k] for k in self.keys())


class BaseDataset(Dataset):
    """This class is the Base class for pythae's dataset

    A ``__getitem__`` is redefined and outputs a python dictionnary
    with the keys corresponding to `data` and `labels`.
    This Class should be used for any new data sets.
    """

    def __init__(self, data, labels):
        self.labels = labels.type(torch.float)
        self.data = data.type(torch.float)

    def __len__(self):
        return len(self.data)

    def __getitem__(self, index):
        """Generates one sample of data

        Args:
            index (int): The index of the data in the Dataset

        Returns:
            (dict): A dictionnary with the keys 'data' and 'labels' and corresponding
            torch.Tensor
        """
        X = self.data[index]
        y = self.labels[index]
        return DatasetOutput(data=X, labels=y)


class BaseDecoder(nn.Module):
    """This is a base class for Decoders neural networks."""

    def __init__(self):
        nn.Module.__init__(self)

    def forward(self, z: torch.Tensor):
        """This function must be implemented in a child class.
        It takes the input data and returns an instance of
        :class:`~pythae.models.base.base_utils.ModelOutput`.
        If you decide to provide your own decoder network, you must make sure your
        model inherit from this class by setting and then defining your forward function as
        such:

        .. code-block::

            >>> from pythae.models.nn import BaseDecoder
            >>> from pythae.models.base.base_utils import ModelOutput
            ...
            >>> class My_decoder(BaseDecoder):
            ...
            ...    def __init__(self):
            ...        BaseDecoder.__init__(self)
            ...        # your code
            ...
            ...    def forward(self, z: torch.Tensor):
            ...        # your code
            ...        output = ModelOutput(
            ...             reconstruction=reconstruction
            ...         )
            ...        return output

        Parameters:
            z (torch.Tensor): The latent data that must be decoded

        Returns:
            output (~pythae.models.base.base_utils.ModelOutput): The output of the decoder

        .. note::

            By convention, the reconstruction tensors should be in [0, 1] and of shape
            BATCH x channels x ...

        """
        raise NotImplementedError()


class BaseEncoder(nn.Module):
    """This is a base class for Encoders neural networks."""

    def __init__(self):
        nn.Module.__init__(self)

    def forward(self, x):
        """This function must be implemented in a child class.
        It takes the input data and returns an instance of
        :class:`~pythae.models.base.base_utils.ModelOutput`.
        If you decide to provide your own encoder network, you must make sure your
        model inherit from this class by setting and then defining your forward function as
        such:

        .. code-block::

            >>> from pythae.models.nn import BaseEncoder
            >>> from pythae.models.base.base_utils import ModelOutput
            ...
            >>> class My_Encoder(BaseEncoder):
            ...
            ...     def __init__(self):
            ...         BaseEncoder.__init__(self)
            ...         # your code
            ...
            ...     def forward(self, x: torch.Tensor):
            ...         # your code
            ...         output = ModelOutput(
            ...             embedding=embedding,
            ...             log_covariance=log_var # for VAE based models
            ...         )
            ...         return output

        Parameters:
            x (torch.Tensor): The input data that must be encoded

        Returns:
            output (~pythae.models.base.base_utils.ModelOutput): The output of the encoder
        """
        raise NotImplementedError()


class ModelOutput(OrderedDict):
    """Base ModelOutput class fixing the output type from the models. This class is inspired from
    the ``ModelOutput`` class from hugginface transformers library"""

    def __getitem__(self, k):
        if isinstance(k, str):
            self_dict = {k: v for k, v in self.items()}
            return self_dict[k]
        else:
            return self.to_tuple()[k]

    def __setattr__(self, name, value):
        super().__setitem__(name, value)
        super().__setattr__(name, value)

    def __setitem__(self, key, value):
        super().__setitem__(key, value)
        super().__setattr__(key, value)

    def to_tuple(self) ->Tuple[Any]:
        """
        Convert self to a tuple containing all the attributes/keys that are not ``None``.
        """
        return tuple(self[k] for k in self.keys())


class Decoder_AE_MLP(BaseDecoder):

    def __init__(self, args: dict):
        BaseDecoder.__init__(self)
        self.input_dim = args.input_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(args.latent_dim, 512), nn.ReLU()))
        layers.append(nn.Sequential(nn.Linear(512, int(np.prod(args.input_dim))), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'reconstruction_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['reconstruction'] = out.reshape((z.shape[0],) + self.input_dim)
        return output


class LoadError(Exception):
    pass


model_card_template = """---
language: en
tags:
- pythae
license: apache-2.0
---

### Downloading this model from the Hub
This model was trained with pythae. It can be downloaded or reloaded using the method `load_from_hf_hub`
```python
>>> from pythae.models import AutoModel
>>> model = AutoModel.load_from_hf_hub(hf_hub_path="your_hf_username/repo_name")
```
"""


class FactorVAEDiscriminator(nn.Module):

    def __init__(self, latent_dim=16, hidden_units=1000) ->None:
        nn.Module.__init__(self)
        self.layers = nn.Sequential(nn.Linear(latent_dim, hidden_units), nn.LeakyReLU(0.2), nn.Linear(hidden_units, hidden_units), nn.LeakyReLU(0.2), nn.Linear(hidden_units, hidden_units), nn.LeakyReLU(0.2), nn.Linear(hidden_units, hidden_units), nn.LeakyReLU(0.2), nn.Linear(hidden_units, hidden_units), nn.LeakyReLU(0.2), nn.Linear(hidden_units, 2))

    def forward(self, z: torch.Tensor):
        return self.layers(z)


class MSSSIM(torch.nn.Module):

    def __init__(self, window_size=11):
        super(MSSSIM, self).__init__()
        self.window_size = window_size

    def _gaussian(self, sigma):
        gauss = torch.Tensor([np.exp(-(x - self.window_size // 2) ** 2 / float(2 * sigma ** 2)) for x in range(self.window_size)])
        return gauss / gauss.sum()

    def _create_window(self, channel=1):
        _1D_window = self._gaussian(1.5).unsqueeze(1)
        _2D_window = _1D_window.mm(_1D_window.t()).float().unsqueeze(0).unsqueeze(0)
        window = _2D_window.expand(channel, 1, self.window_size, self.window_size).contiguous()
        return window

    def ssim(self, img1: torch.Tensor, img2: torch.Tensor):
        padd = int(self.window_size / 2)
        _, channel, height, width = img1.shape
        window = self._create_window(channel=channel)
        mu1 = F.conv2d(img1, window, padding=padd, groups=channel)
        mu2 = F.conv2d(img2, window, padding=padd, groups=channel)
        mu1_sq = mu1.pow(2)
        mu2_sq = mu2.pow(2)
        mu1_mu2 = mu1 * mu2
        sigma1_sq = F.conv2d(img1 * img1, window, padding=padd, groups=channel) - mu1_sq
        sigma2_sq = F.conv2d(img2 * img2, window, padding=padd, groups=channel) - mu2_sq
        sigma12 = F.conv2d(img1 * img2, window, padding=padd, groups=channel) - mu1_mu2
        L = 1
        C1 = (0.01 * L) ** 2
        C2 = (0.03 * L) ** 2
        v1 = 2.0 * sigma12 + C2
        v2 = sigma1_sq + sigma2_sq + C2
        cs = torch.mean(v1 / v2)
        ssim_map = (2 * mu1_mu2 + C1) * v1 / ((mu1_sq + mu2_sq + C1) * v2)
        ret = ssim_map.mean()
        return ret, cs

    def forward(self, img1, img2):
        if img1.shape[-1] < 4:
            weights = torch.FloatTensor([1.0])
        elif img1.shape[-1] < 8:
            weights = torch.FloatTensor([0.3222, 0.6778])
        elif img1.shape[-1] < 16:
            weights = torch.FloatTensor([0.4558, 0.1633, 0.3809])
        elif img1.shape[-1] < 32:
            weights = torch.FloatTensor([0.3117, 0.3384, 0.2675, 0.0824])
        else:
            weights = torch.FloatTensor([0.0448, 0.2856, 0.3001, 0.2363, 0.1333])
        levels = weights.size()[0]
        mssim = []
        mcs = []
        for _ in range(levels):
            sim, cs = self.ssim(img1, img2)
            mssim.append(sim)
            mcs.append(cs)
            img1 = F.avg_pool2d(img1, (2, 2))
            img2 = F.avg_pool2d(img2, (2, 2))
        mssim = torch.stack(mssim)
        mcs = torch.stack(mcs)
        mssim = (mssim + 1) / 2
        mcs = (mcs + 1) / 2
        pow1 = mcs ** weights
        pow2 = mssim ** weights
        output = torch.prod(pow1[:-1] * pow2[-1])
        return 1 - output


class BaseMetric(nn.Module):
    """This is a base class for Metrics neural networks
    (only applicable for Riemannian based VAE)
    """

    def __init__(self):
        nn.Module.__init__(self)

    def forward(self, x):
        """This function must be implemented in a child class.
        It takes the input data and returns an instance of
        :class:`~pythae.models.base.base_utils.ModelOutput`.
        If you decide to provide your own metric network, you must make sure your
        model inherit from this class by setting and then defining your forward function as
        such:

        .. code-block::

            >>> from pythae.models.nn import BaseMetric
            >>> from pythae.models.base.base_utils import ModelOutput
            ...
            >>> class My_Metric(BaseMetric):
            ...
            ...    def __init__(self):
            ...        BaseMetric.__init__(self)
            ...        # your code
            ...
            ...    def forward(self, x: torch.Tensor):
            ...        # your code
            ...        output = ModelOutput(
            ...             L=L # L matrices in the metric of  Riemannian based VAE (see docs)
            ...         )
            ...        return output

        Parameters:
            x (torch.Tensor): The input data that must be encoded

        Returns:
            output (~pythae.models.base.base_utils.ModelOutput): The output of the metric
        """
        raise NotImplementedError()


class BaseDiscriminator(nn.Module):
    """This is a base class for Discriminator neural networks."""

    def __init__(self):
        nn.Module.__init__(self)

    def forward(self, x):
        """This function must be implemented in a child class.
        It takes the input data and returns an instance of
        :class:`~pythae.models.base.base_utils.ModelOutput`.
        If you decide to provide your own disctriminator network, you must make sure your
        model inherit from this class by setting and then defining your forward function as
        such:

        .. code-block::

            >>> from pythae.models.nn import BaseDiscriminator
            >>> from pythae.models.base.base_utils import ModelOutput
            ...
            >>> class My_Discriminator(BaseDiscriminator):
            ...
            ...     def __init__(self):
            ...         BaseDiscriminator.__init__(self)
            ...         # your code
            ...
            ...     def forward(self, x: torch.Tensor):
            ...         # your code
            ...         output = ModelOutput(
            ...             adversarial_cost=adversarial_cost
            ...         )
            ...         return output

        Parameters:
            x (torch.Tensor): The input data that must be encoded

        Returns:
            output (~pythae.models.base.base_utils.ModelOutput): The output of the encoder
        """
        raise NotImplementedError()


class Decoder_Conv_AE_CELEBA(BaseDecoder):
    """
    A Convolutional decoder Neural net suited for CELEBA-64 and Autoencoder-based
    models.

    It can be built as follows:

    .. code-block::

            >>> from pythae.models.nn.benchmarks.celeba import Decoder_Conv_AE_CELEBA
            >>> from pythae.models import VAEConfig
            >>> model_config = VAEConfig(input_dim=(3, 64, 64), latent_dim=64)
            >>> decoder = Decoder_Conv_AE_CELEBA(model_config)
            >>> decoder
            ... Decoder_Conv_AE_CELEBA(
            ...   (layers): ModuleList(
            ...     (0): Sequential(
            ...       (0): Linear(in_features=64, out_features=65536, bias=True)
            ...     )
            ...     (1): Sequential(
            ...       (0): ConvTranspose2d(1024, 512, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2))
            ...       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (2): Sequential(
            ...       (0): ConvTranspose2d(512, 256, kernel_size=(5, 5), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): ConvTranspose2d(256, 128, kernel_size=(5, 5), stride=(2, 2), padding=(2, 2), output_padding=(1, 1))
            ...       (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (4): Sequential(
            ...       (0): ConvTranspose2d(128, 3, kernel_size=(5, 5), stride=(1, 1), padding=(1, 1))
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )


    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAE
        >>> model = VAE(model_config=model_config, decoder=decoder)
        >>> model.decoder == decoder
        ... True

    .. note::

        Please note that this decoder is suitable for **all** models.

        .. code-block::

            >>> import torch
            >>> input = torch.randn(2, 64)
            >>> out = decoder(input)
            >>> out.reconstruction.shape
            ... torch.Size([2, 3, 64, 64])
    """

    def __init__(self, args: dict):
        BaseDecoder.__init__(self)
        self.input_dim = 3, 64, 64
        self.latent_dim = args.latent_dim
        self.n_channels = 3
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(args.latent_dim, 1024 * 8 * 8)))
        layers.append(nn.Sequential(nn.ConvTranspose2d(1024, 512, 5, 2, padding=2), nn.BatchNorm2d(512), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(512, 256, 5, 2, padding=1, output_padding=0), nn.BatchNorm2d(256), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(256, 128, 5, 2, padding=2, output_padding=1), nn.BatchNorm2d(128), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(128, self.n_channels, 5, 1, padding=1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the reconstruction of the latent code
            under the key `reconstruction`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `reconstruction_layer_i`
            where i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z
        for i in range(max_depth):
            out = self.layers[i](out)
            if i == 0:
                out = out.reshape(z.shape[0], 1024, 8, 8)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'reconstruction_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['reconstruction'] = out
        return output


class Discriminator_Conv_CELEBA(BaseDiscriminator):
    """
    A Convolutional discriminator Neural net suited for CELEBA.


    It can be built as follows:

    .. code-block::

            >>> from pythae.models.nn.benchmarks.celeba import Discriminator_Conv_CELEBA
            >>> from pythae.models import VAEGANConfig
            >>> model_config = VAEGANConfig(input_dim=(3, 64, 64), latent_dim=64)
            >>> discriminator = Discriminator_Conv_CELEBA(model_config)
            >>> discriminator
            ... Discriminator_Conv_CELEBA(
            ...   (layers): ModuleList(
            ...     (0): Sequential(
            ...       (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(128, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (1): Sequential(
            ...       (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): Tanh()
            ...     )
            ...     (2): Sequential(
            ...       (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(1024, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (4): Sequential(
            ...       (0): Linear(in_features=16384, out_features=1, bias=True)
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )


    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAEGAN
        >>> model = VAEGAN(model_config=model_config, discriminator=discriminator)
        >>> model.discriminator == discriminator
        ... True
    """

    def __init__(self, args: dict):
        BaseDiscriminator.__init__(self)
        self.input_dim = 3, 64, 64
        self.latent_dim = args.latent_dim
        self.n_channels = 3
        self.discriminator_input_dim = args.discriminator_input_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, 128, 4, 2, padding=1), nn.BatchNorm2d(128), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(128, 256, 4, 2, padding=1), nn.BatchNorm2d(256), nn.Tanh()))
        layers.append(nn.Sequential(nn.Conv2d(256, 512, 4, 2, padding=1), nn.BatchNorm2d(512), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(512, 1024, 4, 2, padding=1), nn.BatchNorm2d(1024), nn.ReLU()))
        layers.append(nn.Sequential(nn.Linear(1024 * 4 * 4, 1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the adversarial score of the input
            under the key `embedding`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `embedding_layer_i` where
            i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            if i == 4:
                out = out.reshape(x.shape[0], -1)
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = out
        return output


class ResBlock(nn.Module):

    def __init__(self, in_channels, out_channels):
        nn.Module.__init__(self)
        self.conv_block = nn.Sequential(nn.ReLU(), nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1), nn.ReLU(), nn.Conv2d(out_channels, in_channels, kernel_size=1, stride=1, padding=0))

    def forward(self, x: torch.tensor) ->torch.Tensor:
        return x + self.conv_block(x)


class Decoder_Conv_AE_CIFAR(BaseDecoder):
    """
    A Convolutional decoder Neural net suited for CIFAR and Autoencoder-based
    models.

    It can be built as follows:

    .. code-block::

            >>> from pythae.models.nn.benchmarks.cifar import Decoder_Conv_AE_CIFAR
            >>> from pythae.models import VAEConfig
            >>> model_config = VAEConfig(input_dim=(3, 32, 32), latent_dim=16)
            >>> decoder = Decoder_Conv_AE_CIFAR(model_config)
            >>> decoder
            ... Decoder_Conv_AE_CIFAR(
            ...   (layers): ModuleList(
            ...     (0): Linear(in_features=16, out_features=65536, bias=True)
            ...     (1): Sequential(
            ...       (0): ConvTranspose2d(1024, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (2): Sequential(
            ...       (0): ConvTranspose2d(512, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
            ...       (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): ConvTranspose2d(256, 3, kernel_size=(4, 4), stride=(1, 1), padding=(2, 2))
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )


    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAE
        >>> model = VAE(model_config=model_config, decoder=decoder)
        >>> model.decoder == decoder
        ... True


    .. note::

        Please note that this decoder is suitable for **all** models.

        .. code-block::

            >>> import torch
            >>> input = torch.randn(2, 16)
            >>> out = decoder(input)
            >>> out.reconstruction.shape
            ... torch.Size([2, 3, 32, 32])
    """

    def __init__(self, args: dict):
        BaseDecoder.__init__(self)
        self.input_dim = 3, 32, 32
        self.latent_dim = args.latent_dim
        self.n_channels = 3
        layers = nn.ModuleList()
        layers.append(nn.Linear(args.latent_dim, 1024 * 8 * 8))
        layers.append(nn.Sequential(nn.ConvTranspose2d(1024, 512, 4, 2, padding=1), nn.BatchNorm2d(512), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(512, 256, 4, 2, padding=1, output_padding=1), nn.BatchNorm2d(256), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(256, self.n_channels, 4, 1, padding=2), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the reconstruction of the latent code
            under the key `reconstruction`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `reconstruction_layer_i`
            where i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z
        for i in range(max_depth):
            out = self.layers[i](out)
            if i == 0:
                out = out.reshape(z.shape[0], 1024, 8, 8)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'reconstruction_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['reconstruction'] = out
        return output


class Discriminator_Conv_CIFAR(BaseDiscriminator):
    """
    A Convolutional discriminator Neural net suited for CIFAR.


    It can be built as follows:

    .. code-block::

            >>> from pythae.models.nn.benchmarks.cifar import Discriminator_Conv_CIFAR
            >>> from pythae.models import VAEGANConfig
            >>> model_config = VAEGANConfig(input_dim=(3, 32, 32), latent_dim=16)
            >>> discriminator = Discriminator_Conv_CIFAR(model_config)
            >>> discriminator
            ... Discriminator_Conv_CIFAR(
            ...   (layers): ModuleList(
            ...     (0): Sequential(
            ...       (0): Conv2d(3, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (1): Sequential(
            ...       (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): Tanh()
            ...     )
            ...     (2): Sequential(
            ...       (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (4): Sequential(
            ...       (0): Linear(in_features=4096, out_features=1, bias=True)
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )

    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAEGAN
        >>> model = VAEGAN(model_config=model_config, discriminator=discriminator)
        >>> model.discriminator == discriminator
        ... True
    """

    def __init__(self, args: dict):
        BaseDiscriminator.__init__(self)
        self.input_dim = 3, 32, 32
        self.latent_dim = args.latent_dim
        self.n_channels = 3
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, 128, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(128, 256, 4, 2, padding=1), nn.Tanh()))
        layers.append(nn.Sequential(nn.Conv2d(256, 512, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(512, 1024, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Linear(1024 * 2 * 2, 1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the adversarial score of the input
            under the key `embedding`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `embedding_layer_i` where
            i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            if i == 4:
                out = out.reshape(x.shape[0], -1)
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = out
        return output


class Decoder_Conv_AE_MNIST(BaseDecoder):
    """
    A Convolutional decoder suited for MNIST and Autoencoder-based
    models.

    .. code-block::

            >>> from pythae.models.nn.benchmarks.mnist import Decoder_Conv_AE_MNIST
            >>> from pythae.models import VAEConfig
            >>> model_config = VAEConfig(input_dim=(1, 28, 28), latent_dim=16)
            >>> decoder = Decoder_Conv_AE_MNIST(model_config)
            >>> decoder
            ... Decoder_Conv_AE_MNIST(
            ...   (layers): ModuleList(
            ...     (0): Linear(in_features=16, out_features=16384, bias=True)
            ...     (1): Sequential(
            ...       (0): ConvTranspose2d(1024, 512, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1))
            ...       (1): BatchNorm2d(512, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (2): Sequential(
            ...       (0): ConvTranspose2d(512, 256, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
            ...       (1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
            ...       (2): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): ConvTranspose2d(256, 1, kernel_size=(3, 3), stride=(2, 2), padding=(1, 1), output_padding=(1, 1))
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )


    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAE
        >>> model = VAE(model_config=model_config, decoder=decoder)
        >>> model.decoder == decoder
        ... True

    .. note::

        Please note that this decoder is suitable for **all** models.

        .. code-block::

            >>> import torch
            >>> input = torch.randn(2, 16)
            >>> out = decoder(input)
            >>> out.reconstruction.shape
            ... torch.Size([2, 1, 28, 28])
    """

    def __init__(self, args: dict):
        BaseDecoder.__init__(self)
        self.input_dim = 1, 28, 28
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Linear(args.latent_dim, 1024 * 4 * 4))
        layers.append(nn.Sequential(nn.ConvTranspose2d(1024, 512, 3, 2, padding=1), nn.BatchNorm2d(512), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(512, 256, 3, 2, padding=1, output_padding=1), nn.BatchNorm2d(256), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(256, self.n_channels, 3, 2, padding=1, output_padding=1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the reconstruction of the latent code
            under the key `reconstruction`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `reconstruction_layer_i`
            where i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}).Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z
        for i in range(max_depth):
            out = self.layers[i](out)
            if i == 0:
                out = out.reshape(z.shape[0], 1024, 4, 4)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'reconstruction_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['reconstruction'] = out
        return output


class Discriminator_Conv_MNIST(BaseDiscriminator):
    """
    A Convolutional discriminator suited for MNIST.


    It can be built as follows:

    .. code-block::

            >>> from pythae.models.nn.benchmarks.mnist import Discriminator_Conv_MNIST
            >>> from pythae.models import VAEGANConfig
            >>> model_config = VAEGANConfig(input_dim=(1, 28, 28), latent_dim=16)
            >>> discriminator = Discriminator_Conv_MNIST(model_config)
            >>> discriminator
            ... Discriminator_Conv_MNIST(
            ...   (layers): ModuleList(
            ...     (0): Sequential(
            ...       (0): Conv2d(1, 128, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (1): Sequential(
            ...       (0): Conv2d(128, 256, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): Tanh()
            ...     )
            ...     (2): Sequential(
            ...       (0): Conv2d(256, 512, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (3): Sequential(
            ...       (0): Conv2d(512, 1024, kernel_size=(4, 4), stride=(2, 2), padding=(1, 1))
            ...       (1): ReLU()
            ...     )
            ...     (4): Sequential(
            ...       (0): Linear(in_features=1024, out_features=1, bias=True)
            ...       (1): Sigmoid()
            ...     )
            ...   )
            ... )

    and then passed to a :class:`pythae.models` instance

        >>> from pythae.models import VAEGAN
        >>> model = VAEGAN(model_config=model_config, discriminator=discriminator)
        >>> model.discriminator == discriminator
        ... True
    """

    def __init__(self, args: dict):
        BaseDiscriminator.__init__(self)
        self.input_dim = 1, 28, 28
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, 128, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(128, 256, 4, 2, padding=1), nn.Tanh()))
        layers.append(nn.Sequential(nn.Conv2d(256, 512, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(512, 1024, 4, 2, padding=1), nn.ReLU()))
        layers.append(nn.Sequential(nn.Linear(1024, 1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Args:
            output_layer_levels (List[int]): The levels of the layers where the outputs are
                extracted. If None, the last layer's output is returned. Default: None.

        Returns:
            ModelOutput: An instance of ModelOutput containing the adversarial score of the input
            under the key `embedding`. Optional: The outputs of the layers specified in
            `output_layer_levels` arguments are available under the keys `embedding_layer_i` where
            i is the layer's level.
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            if i == 4:
                out = out.reshape(x.shape[0], -1)
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = out
        return output


class Encoder_AE_MLP(BaseEncoder):

    def __init__(self, args: dict):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(np.prod(args.input_dim), 512), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)

    def forward(self, x, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x.reshape(-1, np.prod(self.input_dim))
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out)
        return output


class Encoder_VAE_MLP(BaseEncoder):

    def __init__(self, args: dict):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(np.prod(args.input_dim), 512), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)
        self.log_var = nn.Linear(512, self.latent_dim)

    def forward(self, x, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x.reshape(-1, np.prod(self.input_dim))
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out)
                output['log_covariance'] = self.log_var(out)
        return output


class Encoder_SVAE_MLP(BaseEncoder):

    def __init__(self, args: dict):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(np.prod(args.input_dim), 512), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)
        self.log_concentration = nn.Linear(512, 1)

    def forward(self, x, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x.reshape(-1, np.prod(self.input_dim))
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out)
                output['log_concentration'] = self.log_concentration(out)
        return output


class Metric_MLP(BaseMetric):

    def __init__(self, args: dict):
        BaseMetric.__init__(self)
        if args.input_dim is None:
            raise AttributeError("No input dimension provided !'input_dim' parameter of ModelConfig instance must be set to 'data_shape' where the shape of the data is [mini_batch x data_shape]. Unable to build metric automatically")
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.layers = nn.Sequential(nn.Linear(np.prod(args.input_dim), 400), nn.ReLU())
        self.diag = nn.Linear(400, self.latent_dim)
        k = int(self.latent_dim * (self.latent_dim - 1) / 2)
        self.lower = nn.Linear(400, k)

    def forward(self, x):
        h1 = self.layers(x.reshape(-1, np.prod(self.input_dim)))
        h21, h22 = self.diag(h1), self.lower(h1)
        L = torch.zeros((x.shape[0], self.latent_dim, self.latent_dim))
        indices = torch.tril_indices(row=self.latent_dim, col=self.latent_dim, offset=-1)
        L[:, indices[0], indices[1]] = h22
        L = L + torch.diag_embed(h21.exp())
        output = ModelOutput(L=L)
        return output


class Discriminator_MLP(BaseDiscriminator):

    def __init__(self, args: dict):
        BaseDiscriminator.__init__(self)
        self.discriminator_input_dim = args.discriminator_input_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(np.prod(args.discriminator_input_dim), 256), nn.ReLU()))
        layers.append(nn.Sequential(nn.Linear(256, 1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        """Forward method

        Returns:
            ModelOutput: An instance of ModelOutput containing the reconstruction of the latent code
            under the key `reconstruction`
        """
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}). Got ({output_layer_levels}).'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z.reshape(z.shape[0], -1)
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = out
        return output


class BatchNorm(nn.Module):
    """A BatchNorm layer used in several flows"""

    def __init__(self, num_features, eps=1e-05, momentum=0.1):
        nn.Module.__init__(self)
        self.eps = eps
        self.momentum = momentum
        self.log_gamma = nn.Parameter(torch.zeros(num_features))
        self.beta = nn.Parameter(torch.zeros(num_features))
        self.register_buffer('running_mean', torch.zeros(num_features))
        self.register_buffer('running_var', torch.ones(num_features))

    def forward(self, x):
        if self.training:
            self.batch_mean = x.mean(0).data
            self.batch_var = x.var(0).data
            self.running_mean.mul_(1 - self.momentum).add_(self.batch_mean * self.momentum)
            self.running_var.mul_(1 - self.momentum).add_(self.batch_var * self.momentum)
            mean = self.batch_mean
            var = self.batch_var
        else:
            mean = self.running_mean
            var = self.running_var
        y = (x - mean) / (var + self.eps).sqrt() * self.log_gamma.exp() + self.beta
        log_abs_det_jac = self.log_gamma - 0.5 * (var + self.eps).log()
        output = ModelOutput(out=y, log_abs_det_jac=log_abs_det_jac.expand_as(x).sum(dim=-1))
        return output

    def inverse(self, y):
        if self.training:
            if not hasattr(self, 'batch_mean') or not hasattr(self, 'batch_var'):
                mean = torch.zeros(1)
                var = torch.ones(1)
            else:
                mean = self.batch_mean
                var = self.batch_var
        else:
            mean = self.running_mean
            var = self.running_var
        x = (y - self.beta) * (-self.log_gamma).exp() * (var + self.eps).sqrt() + mean
        log_abs_det_jac = -self.log_gamma + 0.5 * (var + self.eps).log()
        output = ModelOutput(out=x, log_abs_det_jac=log_abs_det_jac.expand_as(x).sum(dim=-1))
        return output


class MaskedLinear(nn.Linear):
    """Masked Linear Layer inheriting from `~torch.nn.Linear` class and applying a mask to consider
    only a selection of weight.
    """

    def __init__(self, in_features, out_features, mask):
        nn.Linear.__init__(self, in_features=in_features, out_features=out_features)
        self.register_buffer('mask', mask)

    def forward(self, x):
        return F.linear(x, self.mask * self.weight, self.bias)


class MaskedConv2d(nn.Conv2d):

    def __init__(self, mask_type: str, in_channels: int, out_channels: int, kernel_size: int, stride: int=1, padding: Union[str, int]=0, dilation: int=1, groups: int=1, bias: bool=True, padding_mode: str='zeros', device=None, dtype=None) ->None:
        super().__init__(in_channels, out_channels, kernel_size, stride, padding, dilation, groups, bias, padding_mode, device, dtype)
        self.register_buffer('mask', torch.ones_like(self.weight))
        _, _, kH, kW = self.weight.shape
        if mask_type == 'A':
            self.mask[:, :, kH // 2, kW // 2:] = 0
            self.mask[:, :, kH // 2 + 1:] = 0
        else:
            self.mask[:, :, kH // 2, kW // 2 + 1:] = 0
            self.mask[:, :, kH // 2 + 1:] = 0

    def forward(self, input: torch.Tensor) ->torch.Tensor:
        return F.conv2d(input, self.weight * self.mask, self.bias, self.stride, self.padding, self.dilation)


ACTIVATION = {'elu': F.elu, 'tanh': torch.tanh, 'linear': lambda x: x}


ACTIVATION_DERIVATIVES = {'elu': lambda x: torch.ones_like(x) * (x >= 0) + torch.exp(x) * (x < 0), 'tanh': lambda x: 1 - torch.tanh(x) ** 2, 'linear': lambda x: 1}


class Layer(nn.Module):

    def __init__(self) ->None:
        nn.Module.__init__(self)

    def forward(self, x):
        return x


class Encoder_AE_Conv(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}) indice. Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out.reshape(x.shape[0], -1))
        return output


class Encoder_VAE_Conv(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)
        self.log_var = nn.Linear(512, self.latent_dim)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}) indice. Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out.reshape(x.shape[0], -1))
                output['log_covariance'] = self.log_var(out.reshape(x.shape[0], -1))
        return output


class Encoder_SVAE_Conv(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Conv2d(self.n_channels, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.Conv2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        self.layers = layers
        self.depth = len(layers)
        self.embedding = nn.Linear(512, self.latent_dim)
        self.log_concentration = nn.Linear(512, 1)

    def forward(self, x: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}) indice. Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = x
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = self.embedding(out.reshape(x.shape[0], -1))
                output['log_concentration'] = self.log_concentration(out.reshape(x.shape[0], -1))
        return output


class Decoder_AE_Conv(BaseDecoder):

    def __init__(self, args):
        BaseDecoder.__init__(self)
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.n_channels = 1
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(self.latent_dim, 256), nn.ReLU(), nn.Linear(256, 512), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(32, out_channels=32, kernel_size=3, stride=2, padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(32, out_channels=32, kernel_size=3, stride=2, padding=1, output_padding=1), nn.BatchNorm2d(32), nn.ReLU()))
        layers.append(nn.Sequential(nn.ConvTranspose2d(32, out_channels=self.n_channels, kernel_size=3, stride=2, padding=1, output_padding=1), nn.BatchNorm2d(self.n_channels), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}) indice. Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'reconstruction_layer_{i + 1}'] = out
            if i == 0:
                out = out.reshape(z.shape[0], 32, 4, 4)
            if i + 1 == self.depth:
                output['reconstruction'] = out
        return output


class Metric_Custom(BaseMetric):

    def __init__(self):
        BaseMetric.__init__(self)


class Encoder_AE_MLP_Custom(BaseEncoder):

    def __init__(self, args: dict):
        BaseEncoder.__init__(self)
        if args.input_dim is None:
            raise AttributeError("No input dimension provided !'input_dim' parameter of ModelConfig instance must be set to 'data_shape' wherethe shape of the data is [mini_batch x data_shape]. Unable to build encoderautomatically")
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.layers = nn.Sequential(nn.Linear(np.prod(args.input_dim), 10), nn.ReLU(), Layer())
        self.mu = nn.Linear(10, self.latent_dim)

    def forward(self, x):
        out = self.layers(x.reshape(-1, int(np.prod(self.input_dim))))
        output = ModelOutput(embedding=self.mu(out))
        return output


class Encoder_VAE_MLP_Custom(BaseEncoder):

    def __init__(self, args: dict):
        BaseEncoder.__init__(self)
        if args.input_dim is None:
            raise AttributeError("No input dimension provided !'input_dim' parameter of ModelConfig instance must be set to 'data_shape' wherethe shape of the data is [mini_batch x data_shape]. Unable to build encoderautomatically")
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.layers = nn.Sequential(nn.Linear(np.prod(args.input_dim), 10), nn.ReLU(), Layer())
        self.mu = nn.Linear(10, self.latent_dim)
        self.std = nn.Linear(10, self.latent_dim)

    def forward(self, x):
        out = self.layers(x.reshape(-1, int(np.prod(self.input_dim))))
        output = ModelOutput(embedding=self.mu(out), log_covariance=self.std(out))
        return output


class Decoder_MLP_Custom(BaseDecoder):

    def __init__(self, args: dict):
        BaseDecoder.__init__(self)
        if args.input_dim is None:
            raise AttributeError("No input dimension provided !'input_dim' parameter of ModelConfig instance must be set to 'data_shape' wherethe shape of the data is [mini_batch x data_shape]. Unable to build decoderautomatically")
        self.layers = nn.Sequential(nn.Linear(args.latent_dim, 10), nn.ReLU(), nn.Linear(10, np.prod(args.input_dim)), nn.Sigmoid())

    def forward(self, z):
        out = self.layers(z)
        output = ModelOutput(reconstruction=out)
        return output


class Metric_MLP_Custom(BaseMetric):

    def __init__(self, args: dict):
        BaseMetric.__init__(self)
        if args.input_dim is None:
            raise AttributeError("No input dimension provided !'input_dim' parameter of ModelConfig instance must be set to 'data_shape' wherethe shape of the data is [mini_batch x data_shape]. Unable to build metricautomatically")
        self.input_dim = args.input_dim
        self.latent_dim = args.latent_dim
        self.layers = nn.Sequential(nn.Linear(np.prod(self.input_dim), 10), nn.ReLU())
        self.diag = nn.Linear(10, self.latent_dim)
        k = int(self.latent_dim * (self.latent_dim - 1) / 2)
        self.lower = nn.Linear(10, k)

    def forward(self, x):
        h1 = self.layers(x.reshape(-1, int(np.prod(self.input_dim))))
        h21, h22 = self.diag(h1), self.lower(h1)
        L = torch.zeros((x.shape[0], self.latent_dim, self.latent_dim))
        indices = torch.tril_indices(row=self.latent_dim, col=self.latent_dim, offset=-1)
        L[:, indices[0], indices[1]] = h22
        L = L + torch.diag_embed(h21.exp())
        output = ModelOutput(L=L)
        return output


class Discriminator_MLP_Custom(BaseDiscriminator):

    def __init__(self, args: dict):
        BaseDiscriminator.__init__(self)
        self.discriminator_input_dim = args.discriminator_input_dim
        layers = nn.ModuleList()
        layers.append(nn.Sequential(nn.Linear(np.prod(args.discriminator_input_dim), 10), nn.ReLU(inplace=True)))
        layers.append(nn.Linear(10, 5))
        layers.append(nn.Sequential(nn.Linear(5, 1), nn.Sigmoid()))
        self.layers = layers
        self.depth = len(layers)

    def forward(self, z: torch.Tensor, output_layer_levels: List[int]=None):
        output = ModelOutput()
        max_depth = self.depth
        if output_layer_levels is not None:
            assert all(self.depth >= levels > 0 or levels == -1 for levels in output_layer_levels), f'Cannot output layer deeper than depth ({self.depth}) indice. Got ({output_layer_levels})'
            if -1 in output_layer_levels:
                max_depth = self.depth
            else:
                max_depth = max(output_layer_levels)
        out = z.reshape(z.shape[0], -1)
        for i in range(max_depth):
            out = self.layers[i](out)
            if output_layer_levels is not None:
                if i + 1 in output_layer_levels:
                    output[f'embedding_layer_{i + 1}'] = out
            if i + 1 == self.depth:
                output['embedding'] = out
        return output


class EncoderWrongInputDim(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)) - 1, args.latent_dim)

    def forward(self, x):
        output = ModelOutput(embedding=self.fc(x.reshape(-1, int(np.prod(self.input_dim)))))
        return output


class DecoderWrongInputDim(BaseDecoder):

    def __init__(self, args):
        BaseDecoder.__init__(self)
        self.latent_dim = args.latent_dim
        self.fc = nn.Linear(args.latent_dim - 1, int(np.prod(args.input_dim)))

    def forward(self, z):
        out = self.fc(z.reshape(-1, self.latent_dim))
        output = ModelOutput(reconstruction=out)
        return output


class MetricWrongInputDim(BaseMetric):

    def __init__(self, args):
        BaseMetric.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)) - 1, args.latent_dim)

    def forward(self, x):
        L = self.fc(x.reshape(-1, int(np.prod(self.input_dim))))
        output = ModelOutput(L=L)
        return output


class EncoderWrongOutputDim(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)), args.latent_dim - 1)

    def forward(self, x):
        output = ModelOutput(embedding=self.fc(x.reshape(-1, self.input_dim)))
        return output


class DecoderWrongOutputDim(BaseDecoder):

    def __init__(self, args):
        BaseDecoder.__init__(self)
        self.latent_dim = args.latent_dim
        self.fc = nn.Linear(args.latent_dim, int(np.prod(args.input_dim)) - 1)

    def forward(self, z):
        out = self.fc(z.reshape(-1, self.latent_dim))
        output = ModelOutput(reconstruction=out)
        return output


class MetricWrongOutputDim(BaseMetric):

    def __init__(self, args):
        BaseMetric.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)), args.latent_dim - 1)

    def forward(self, x):
        L = self.fc(x.reshape(-1, int(np.prod(self.input_dim))))
        output = ModelOutput(L=L)
        return output


class EncoderWrongOutput(BaseEncoder):

    def __init__(self, args):
        BaseEncoder.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)), args.latent_dim)

    def forward(self, x):
        output = ModelOutput(embedding=self.fc(x.reshape(-1, int(np.prod(self.input_dim)))))
        return output


class DecoderWrongOutput(BaseDecoder):

    def __init__(self, args):
        BaseDecoder.__init__(self)
        self.latent_dim = args.latent_dim
        self.fc = nn.Linear(args.latent_dim, int(np.prod(args.input_dim)))

    def forward(self, z):
        out = self.fc(z.reshape(-1, self.latent_dim))
        output = ModelOutput(reconstruction=out)
        return output, output, output


class MetricWrongOutput(BaseMetric):

    def __init__(self, args):
        BaseMetric.__init__(self)
        self.input_dim = args.input_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)), args.latent_dim)

    def forward(self, x):
        L = self.fc(x.reshape(-1, self.input_dim))
        output = ModelOutput(L=L)
        return output, output


class MetricWrongOutputDimBis(BaseMetric):

    def __init__(self, args):
        BaseMetric.__init__(self)
        self.latent_dim = args.latent_dim
        self.fc = nn.Linear(int(np.prod(args.input_dim)), args.latent_dim)

    def forward(self, x):
        return torch.randn(x.shape[0], self.latent_dim, self.latent_dim - 1)


class NetBadInheritance(nn.Module):

    def __init__(self):
        nn.Module.__init__(self)

    def forward(self, x):
        return 0


import torch
from torch.nn import MSELoss, ReLU
from _paritybench_helpers import _mock_config, _mock_layer, _paritybench_base, _fails_compile


TESTCASES = [
    # (nn.Module, init_args, forward_args, jit_compiles)
    (BatchNorm,
     lambda: ([], {'num_features': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (DecoderWrongOutput,
     lambda: ([], {'args': _mock_config(latent_dim=4, input_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (DecoderWrongOutputDim,
     lambda: ([], {'args': _mock_config(latent_dim=4, input_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Decoder_AE_Conv,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Decoder_Conv_AE_CELEBA,
     lambda: ([], {'args': _mock_config(latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Decoder_Conv_AE_CIFAR,
     lambda: ([], {'args': _mock_config(latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Decoder_Conv_AE_MNIST,
     lambda: ([], {'args': _mock_config(latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Decoder_MLP_Custom,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Discriminator_Conv_CELEBA,
     lambda: ([], {'args': _mock_config(latent_dim=4, discriminator_input_dim=4)}),
     lambda: ([torch.rand([4, 3, 64, 64])], {}),
     False),
    (Discriminator_MLP,
     lambda: ([], {'args': _mock_config(discriminator_input_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Discriminator_MLP_Custom,
     lambda: ([], {'args': _mock_config(discriminator_input_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (EncoderWrongOutput,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (EncoderWrongOutputDim,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Encoder_AE_MLP,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Encoder_AE_MLP_Custom,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Encoder_SVAE_MLP,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Encoder_VAE_MLP,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (Encoder_VAE_MLP_Custom,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (GatedDense,
     lambda: ([], {'input_size': 4, 'output_size': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Layer,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MSSSIM,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4]), torch.rand([4, 4, 4, 4])], {}),
     False),
    (MaskedConv2d,
     lambda: ([], {'mask_type': 4, 'in_channels': 4, 'out_channels': 4, 'kernel_size': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (MetricWrongOutput,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (MetricWrongOutputDim,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
    (MetricWrongOutputDimBis,
     lambda: ([], {'args': _mock_config(latent_dim=4, input_dim=4)}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (Metric_MLP,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (Metric_MLP_Custom,
     lambda: ([], {'args': _mock_config(input_dim=4, latent_dim=4)}),
     lambda: ([torch.rand([4, 4])], {}),
     False),
    (NetBadInheritance,
     lambda: ([], {}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (NonLinear,
     lambda: ([], {'input_size': 4, 'output_size': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     True),
    (ResBlock,
     lambda: ([], {'in_channels': 4, 'out_channels': 4}),
     lambda: ([torch.rand([4, 4, 4, 4])], {}),
     False),
]

class Test_clementchadebec_benchmark_VAE(_paritybench_base):
    def test_000(self):
        self._check(*TESTCASES[0])

    def test_001(self):
        self._check(*TESTCASES[1])

    def test_002(self):
        self._check(*TESTCASES[2])

    def test_003(self):
        self._check(*TESTCASES[3])

    def test_004(self):
        self._check(*TESTCASES[4])

    def test_005(self):
        self._check(*TESTCASES[5])

    def test_006(self):
        self._check(*TESTCASES[6])

    def test_007(self):
        self._check(*TESTCASES[7])

    def test_008(self):
        self._check(*TESTCASES[8])

    def test_009(self):
        self._check(*TESTCASES[9])

    def test_010(self):
        self._check(*TESTCASES[10])

    def test_011(self):
        self._check(*TESTCASES[11])

    def test_012(self):
        self._check(*TESTCASES[12])

    def test_013(self):
        self._check(*TESTCASES[13])

    def test_014(self):
        self._check(*TESTCASES[14])

    def test_015(self):
        self._check(*TESTCASES[15])

    def test_016(self):
        self._check(*TESTCASES[16])

    def test_017(self):
        self._check(*TESTCASES[17])

    def test_018(self):
        self._check(*TESTCASES[18])

    def test_019(self):
        self._check(*TESTCASES[19])

    def test_020(self):
        self._check(*TESTCASES[20])

    def test_021(self):
        self._check(*TESTCASES[21])

    def test_022(self):
        self._check(*TESTCASES[22])

    def test_023(self):
        self._check(*TESTCASES[23])

    def test_024(self):
        self._check(*TESTCASES[24])

    def test_025(self):
        self._check(*TESTCASES[25])

    def test_026(self):
        self._check(*TESTCASES[26])

    def test_027(self):
        self._check(*TESTCASES[27])

    def test_028(self):
        self._check(*TESTCASES[28])

    def test_029(self):
        self._check(*TESTCASES[29])

