class_docstr: Combination Importance Weighted Autoencoder model.
  
  Args:
      model_config (CIWAEConfig): The CIWAE 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, mu, log_var, z):
      if self.model_config.reconstruction_loss == "mse":
          recon_loss = 0.5 * F.mse_loss(
              recon_x,
              x.reshape(recon_x.shape[0], -1)
              .unsqueeze(1)
              .repeat(1, self.n_samples, 1),
              reduction="none",
          ).sum(dim=-1)
  
      elif self.model_config.reconstruction_loss == "bce":
          recon_loss = F.binary_cross_entropy(
              recon_x,
              x.reshape(recon_x.shape[0], -1)
              .unsqueeze(1)
              .repeat(1, self.n_samples, 1),
              reduction="none",
          ).sum(dim=-1)
  
      log_q_z = (
          -0.5 * (log_var + torch.pow(z - mu, 2) / log_var.exp())
      ).sum(dim=-1)
      log_p_z = -0.5 * (z**2).sum(dim=-1)
  
      KLD = -(log_p_z - log_q_z)
  
      log_w = -(recon_loss + KLD)
  
      log_w_minus_max = log_w - log_w.max(1, keepdim=True)[0]
      w = log_w_minus_max.exp()
      w_tilde = (w / w.sum(axis=1, keepdim=True)).detach()
  
      iwae_elbo = -(w_tilde * log_w).sum(1)
      vae_elbo = -log_w.mean(dim=-1)  # avg on importance samples
  
      return (
          (self.beta * vae_elbo + (1 - self.beta) * iwae_elbo).mean(
              dim=0
          ),
          recon_loss.mean(dim=0),
          KLD.mean(dim=0),
      )
  
file_path: benchmark_VAE/src/pythae/models/ciwae/ciwae_model.py
imports: 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 .ciwae_config import CIWAEConfig
label_desc: Use image node_id 3 for calling the CIWAE obj's underlying member method code with example usage: obj.loss_function(recon_x, x, mu, log_var, z) and returns: unknown,  recon_loss,  KLD