class_docstr: Variational Autoencoder using Adversarial reconstruction loss model.
  
  Args:
      model_config (VAEGANConfig): The Autoencoder configuration setting the main
          parameters of the model.
  
      encoder (BaseEncoder): An instance of BaseEncoder (inheriting from `torch.nn.Module` which
          plays the role of encoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
      decoder (BaseDecoder): An instance of BaseDecoder (inheriting from `torch.nn.Module` which
          plays the role of decoder. This argument allows you to use your own neural networks
          architectures if desired. If None is provided, a simple Multi Layer Preception
          (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default: None.
  
      discriminator (BaseDiscriminator): An instance of BaseDiscriminator (inheriting from
          `torch.nn.Module` which plays the role of discriminator. This argument allows you to use
          your own neural networks architectures if desired. If None is provided, a simple Multi
          Layer Preception (https://en.wikipedia.org/wiki/Multilayer_perceptron) is used. Default:
          None.
  
  .. note::
      For high dimensional data we advice you to provide you own network architectures. With the
      provided MLP you may end up with a ``MemoryError``.
class_docstr_tok: ["Variational","Autoencoder","using","Adversarial","reconstruction","loss","model",".","Args",":","model_config","(","VAEGANConfig",")",":","The","Autoencoder","configuration","setting","the","main","parameters","of","the","model",".","encoder","(","BaseEncoder",")",":","An","instance","of","BaseEncoder","(","inheriting","from","`","torch.nn.Module","`","which","plays","the","role","of","encoder",".","This","argument","allows","you","to","use","your","own","neural","networks","architectures","if","desired",".","If","None","is","provided",",","a","simple","Multi","Layer","Preception","(","https",":","\/\/en.wikipedia.org\/wiki\/Multilayer_perceptron",")","is","used",".","Default",":","None",".","decoder","(","BaseDecoder",")",":","An","instance","of","BaseDecoder","(","inheriting","from","`","torch.nn.Module","`","which","plays","the","role","of","decoder",".","This","argument","allows","you","to","use","your","own","neural","networks","architectures","if","desired",".","If","None","is","provided",",","a","simple","Multi","Layer","Preception","(","https",":","\/\/en.wikipedia.org\/wiki\/Multilayer_perceptron",")","is","used",".","Default",":","None",".","discriminator","(","BaseDiscriminator",")",":","An","instance","of","BaseDiscriminator","(","inheriting","from","`","torch.nn.Module","`","which","plays","the","role","of","discriminator",".","This","argument","allows","you","to","use","your","own","neural","networks","architectures","if","desired",".","If","None","is","provided",",","a","simple","Multi","Layer","Preception","(","https",":","\/\/en.wikipedia.org\/wiki\/Multilayer_perceptron",")","is","used",".","Default",":","None",".","..","note",":",":","For","high","dimensional","data","we","advice","you","to","provide","you","own","network","architectures",".","With","the","provided","MLP","you","may","end","up","with","a","``","MemoryError","``","."]
code: def loss_function(self, recon_x, x, z, z_prior, mu, log_var):
      N = z.shape[0]  # batch size
  
      # KL between prior and posterior
      KLD = -0.5 * torch.sum(
          1 + log_var - mu.pow(2) - log_var.exp(), dim=-1
      )
  
      # feature maps of true data
      true_discr_layer = self.discriminator(
          x, output_layer_levels=[self.reconstruction_layer]
      )[f"embedding_layer_{self.reconstruction_layer}"]
  
      # feature maps of recon data
      recon_discr_layer = self.discriminator(
          recon_x, output_layer_levels=[self.reconstruction_layer]
      )[f"embedding_layer_{self.reconstruction_layer}"]
  
      # MSE in feature space
      recon_loss = 0.5 * F.mse_loss(
          true_discr_layer.reshape(N, -1),
          recon_discr_layer.reshape(N, -1),
          reduction="none",
      ).sum(dim=-1)
  
      encoder_loss = KLD + recon_loss
  
      gen_prior = self.decoder(z_prior).reconstruction
  
      true_adversarial_score = self.discriminator(x).embedding.flatten()
      # gen_adversarial_score = self.discriminator(recon_x).embedding.flatten()
      prior_adversarial_score = self.discriminator(
          gen_prior
      ).embedding.flatten()
  
      true_labels = torch.ones(N, requires_grad=False).to(x.device)
      fake_labels = torch.zeros(N, requires_grad=False).to(x.device)
  
      original_dis_cost = F.binary_cross_entropy(
          true_adversarial_score, true_labels
      )  # original are true
      prior_dis_cost = F.binary_cross_entropy(
          prior_adversarial_score, fake_labels
      )  # prior is false
      # gen_cost =  F.binary_cross_entropy(
      #   gen_adversarial_score, fake_labels
      # ) # generated are false
  
      discriminator_loss = (
          (original_dis_cost)
          + (prior_dis_cost)
          # +
          # (gen_cost)
      )
  
      decoder_loss = (
          (1 - self.adversarial_loss_scale) * recon_loss
          - self.adversarial_loss_scale * discriminator_loss
      )
  
      update_encoder = True
      update_discriminator = True
      update_decoder = True
  
      # margins for training stability
      if (
          original_dis_cost.mean() < self.equilibrium - self.margin
          or prior_dis_cost.mean() < self.equilibrium - self.margin
      ):
          update_discriminator = False
  
      if (
          original_dis_cost.mean() > self.equilibrium + self.margin
          or prior_dis_cost.mean() > self.equilibrium + self.margin
      ):
          update_decoder = False
  
      if not update_decoder and not update_discriminator:
          update_discriminator = True
          update_decoder = True
  
      return (
          (recon_loss).mean(dim=0),
          (encoder_loss).mean(dim=0),
          (decoder_loss).mean(dim=0),
          (discriminator_loss).mean(dim=0),
          update_encoder,
          update_decoder,
          update_discriminator,
      )
  
imports: import inspect
  import logging
  import os
  import warnings
  from copy import deepcopy
  from typing import Optional
  import cloudpickle
  import torch
  import torch.nn.functional
  from ...customexception import BadInheritanceError
  from ...data.datasets import BaseDataset
  from ..base.base_utils import CPU_Unpickler, ModelOutput, hf_hub_is_available
  from ..nn import BaseDecoder, BaseDiscriminator, BaseEncoder
  from ..nn.default_architectures import Discriminator_MLP
  from ..vae import VAE
  from .vae_gan_config import VAEGANConfig
label_desc: Use image node_id 4 for calling the VAEGAN obj's underlying member method code with example usage: obj.loss_function(recon_x, x, z, z_prior, mu, log_var) and returns: recon_loss,  encoder_loss,  decoder_loss,  discriminator_loss,  update_encoder,  update_decoder,  update_discriminator