class_docstr: Info Variational Autoencoder model.
  
  Args:
      model_config (INFOVAE_MMD_Config): 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.
  
  .. 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: ["Info","Variational","Autoencoder","model",".","Args",":","model_config","(","INFOVAE_MMD_Config",")",":","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",".","..","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
  
      if self.model_config.reconstruction_loss == "mse":
          recon_loss = 0.5 * F.mse_loss(
              recon_x.reshape(x.shape[0], -1),
              x.reshape(x.shape[0], -1),
              reduction="none",
          ).sum(dim=-1)
  
      elif self.model_config.reconstruction_loss == "bce":
          recon_loss = F.binary_cross_entropy(
              recon_x.reshape(x.shape[0], -1),
              x.reshape(x.shape[0], -1),
              reduction="none",
          ).sum(dim=-1)
  
      KLD = -0.5 * torch.sum(
          1 + log_var - mu.pow(2) - log_var.exp(), dim=-1
      )
  
      if self.kernel_choice == "rbf":
          k_z = self.rbf_kernel(z, z)
          k_z_prior = self.rbf_kernel(z_prior, z_prior)
          k_cross = self.rbf_kernel(z, z_prior)
  
      else:
          k_z = self.imq_kernel(z, z)
          k_z_prior = self.imq_kernel(z_prior, z_prior)
          k_cross = self.imq_kernel(z, z_prior)
  
      mmd_z = (k_z - k_z.diag().diag()).sum() / ((N - 1) * N)
      mmd_z_prior = (k_z_prior - k_z_prior.diag().diag()).sum() / (
          (N - 1) * N
      )
      mmd_cross = k_cross.sum() / (N**2)
  
      mmd_loss = mmd_z + mmd_z_prior - 2 * mmd_cross
  
      loss = (
          recon_loss
          + (1 - self.alpha) * KLD
          + (self.alpha + self.lbd - 1) * mmd_loss
      )
  
      return (
          (loss).mean(dim=0),
          (recon_loss).mean(dim=0),
          (KLD).mean(dim=0),
          mmd_loss,
      )
  
file_path: benchmark_VAE/src/pythae/models/info_vae/info_vae_model.py
imports: import os
  from typing import Optional
  import torch
  import torch.nn.functional
  from ...data.datasets import BaseDataset
  from ..base.base_utils import ModelOutput
  from ..nn import BaseDecoder, BaseEncoder
  from ..vae import VAE
  from .info_vae_config import INFOVAE_MMD_Config
label_desc: Use image node_id 3 for calling the INFOVAE_MMD obj's underlying member method code with example usage: obj.loss_function(recon_x, x, z, z_prior, mu, log_var) and returns: loss,  recon_loss,  KLD,  mmd_loss