class_docstr: :math:`\beta`-TCVAE model.
  
  Args:
      model_config (BetaTCVAEConfig): The Variational 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: [":","math",":","`","\\beta","`","-TCVAE","model",".","Args",":","model_config","(","BetaTCVAEConfig",")",":","The","Variational","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, mu, log_var, z, dataset_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)
  
      log_q_z_given_x = self._compute_log_gauss_density(
          z, mu, log_var
      ).sum(
          dim=-1
      )  # [B]
  
      log_prior = self._compute_log_gauss_density(
          z, torch.zeros_like(z), torch.zeros_like(z)
      ).sum(
          dim=-1
      )  # [B]
  
      log_q_batch_perm = self._compute_log_gauss_density(
          z.reshape(z.shape[0], 1, -1),
          mu.reshape(1, z.shape[0], -1),
          log_var.reshape(1, z.shape[0], -1),
      )  # [B x B x Latent_dim]
  
      if self.use_mss:
          logiw_mat = self._log_importance_weight_matrix(
              z.shape[0], dataset_size
          ).to(z.device)
          log_q_z = torch.logsumexp(
              logiw_mat + log_q_batch_perm.sum(dim=-1), dim=-1
          )  # MMS [B]
          log_prod_q_z = (
              torch.logsumexp(
                  logiw_mat.reshape(z.shape[0], z.shape[0], -1)
                  + log_q_batch_perm,
                  dim=1,
              )
          ).sum(
              dim=-1
          )  # MMS [B]
  
      else:
          log_q_z = torch.logsumexp(
              log_q_batch_perm.sum(dim=-1), dim=-1
          ) - torch.log(
              torch.tensor([z.shape[0] * dataset_size]).to(z.device)
          )  # MWS [B]
          log_prod_q_z = (
              torch.logsumexp(log_q_batch_perm, dim=1)
              - torch.log(
                  torch.tensor([z.shape[0] * dataset_size]).to(z.device)
              )
          ).sum(
              dim=-1
          )  # MWS [B]
  
      mutual_info_loss = log_q_z_given_x - log_q_z
      TC_loss = log_q_z - log_prod_q_z
      dimension_wise_KL = log_prod_q_z - log_prior
  
      return (
          (
              recon_loss
              + self.alpha * mutual_info_loss
              + self.beta * TC_loss
              + self.gamma * dimension_wise_KL
          ).mean(dim=0),
          recon_loss.mean(dim=0),
          (
              self.alpha * mutual_info_loss
              + self.beta * TC_loss
              + self.gamma * dimension_wise_KL
          ).mean(dim=0),
      )
  
file_path: benchmark_VAE/src/pythae/models/beta_tc_vae/beta_tc_vae_model.py
imports: import os
  from typing import Optional
  import numpy
  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 .beta_tc_vae_config import BetaTCVAEConfig
label_desc: Use image node_id 3 for calling the BetaTCVAE obj's underlying member method code with example usage: obj.loss_function(recon_x, x, mu, log_var, z, dataset_size) and returns: unknown,  recon_loss,  unknown