class_docstr: Machine translation model which supports bottleneck architecture,
  NLL, VAE, and MIM loss.
  
  Supported losses:
    1) nll - Conditional cross entropy (the usual NMT loss)
    2) mim - MIM learning framework. A latent variable model with good
                    reconstruction and compressed latent representation.
                    https://arxiv.org/pdf/2003.02645.pdf
    3) vae - VAE learning framework. A latent variable model which learns
                    good probability estimation over observations and
                    a regularized latent representation.
                    https://arxiv.org/pdf/1312.6114.pdf
class_docstr_tok: ["Machine","translation","model","which","supports","bottleneck","architecture",",","NLL",",","VAE",",","and","MIM","loss",".","Supported","losses",":","1",")","nll","-","Conditional","cross","entropy","(","the","usual","NMT","loss",")","2",")","mim","-","MIM","learning","framework",".","A","latent","variable","model","with","good","reconstruction","and","compressed","latent","representation",".","https",":","\/\/arxiv.org\/pdf\/2003.02645.pdf","3",")","vae","-","VAE","learning","framework",".","A","latent","variable","model","which","learns","good","probability","estimation","over","observations","and","a","regularized","latent","representation",".","https",":","\/\/arxiv.org\/pdf\/1312.6114.pdf"]
docstr: Compute the loss from latent (z) and target (x).
  
  train - If True enables loss annealing, and label smoothing
docstr_tok: ["Compute","the","loss","from","latent","(","z",")","and","target","(","x",")",".","train","-","If","True","enables","loss","annealing",",","and","label","smoothing"]
code: def loss(
      self,
      z,
      z_mean,
      z_logv,
      z_mask,
      tgt_log_probs,
      tgt,
      tgt_mask,
      tgt_labels,
      train=False,
      return_info=False,
  ):
      """
      Compute the loss from latent (z) and target (x).
  
      train - If True enables loss annealing, and label smoothing
      """
  
      recon_loss_fn = self.loss_fn if train else self.eval_loss_fn
  
      info_dict = {}
  
      if self.recon_per_token:
          log_p_x_given_z_per_token = -recon_loss_fn(
              log_probs=tgt_log_probs, labels=tgt_labels
          )
  
          log_p_x_given_z = log_p_x_given_z_per_token
          log_p_x_given_z_per_token = log_p_x_given_z_per_token.detach()
      else:
          # averaging of log_p_x_given_z per sample
          output_mask = (
              tgt_labels != self.decoder_tokenizer.pad_id
          ).type_as(tgt_log_probs)
  
          log_p_x_given_z_per_token = (
              -recon_loss_fn(
                  log_probs=tgt_log_probs,
                  labels=tgt_labels,
              ).view(tgt_log_probs.shape[:2])
              * output_mask
          )
  
          # probability per sample
          log_p_x_given_z = log_p_x_given_z_per_token.sum(-1).mean()
  
          tokens = output_mask.sum()
          log_p_x_given_z_per_token = (
              log_p_x_given_z_per_token.sum().detach() / tokens
          )
  
          info_dict["log_p_x_given_z"] = log_p_x_given_z.detach().cpu()
  
      info_dict[
          "log_p_x_given_z_per_token"
      ] = log_p_x_given_z_per_token.detach().cpu()
  
      # loss warmup during training only
      if train:
          trainer = self.trainer
          # if we do not have a trainer ignore annealing
          if trainer is None:
              # ignore warmup and auxiliary loss
              warmup_coef = 1.0
          else:
              global_step = self.trainer.global_step
  
              warmup_coef = min(
                  global_step / self.non_recon_warmup_batches, 1
              )
      else:
          # ignore warmup and auxiliary loss
          warmup_coef = 1.0
  
      info_dict["warmup_coef_recon"] = warmup_coef
  
      if self.model_type in ["mim", "vae"]:
          # tokens = tgt_mask.sum()
          q_z_given_x = torch.distributions.Normal(
              loc=z_mean,
              scale=torch.exp(0.5 * z_logv),
          )
          # average latent distribution to match averaging of observations
          if self.recon_per_token:
              # average latent per dimension - to heuristically match per-token reconstruction
              log_q_z_given_x = (
                  q_z_given_x.log_prob(z).mean(-1).mean(-1).mean()
              )
          else:
              log_q_z_given_x = (
                  q_z_given_x.log_prob(z).sum(-1).sum(-1).mean()
              )
  
          # build prior distribution
          p_z = torch.distributions.Normal(
              loc=torch.zeros_like(z),
              scale=torch.ones_like(z),
          )
          if self.recon_per_token:
              # average latent distribution similar to averaging of observations
              log_p_z = p_z.log_prob(z).mean(-1).mean(-1).mean()
          else:
              log_p_z = p_z.log_prob(z).sum(-1).sum(-1).mean()
  
          if self.model_type == "mim":
              loss_terms = 0.5 * (log_q_z_given_x + log_p_z)
          elif self.model_type == "vae":
              # KL divergence -Dkl( q(z|x) || p(z) )
              loss_terms = log_p_z - log_q_z_given_x
  
          # show loss value for reconstruction but train with MIM/VAE loss
          loss = -(log_p_x_given_z + warmup_coef * loss_terms)
  
          info_dict["log_q_z_given_x"] = log_q_z_given_x.detach().cpu()
          info_dict["log_var_q_z_given_x"] = (
              z_logv.detach().mean().cpu()
          )
          info_dict["log_p_z"] = log_p_z.detach().cpu()
          info_dict["kl_div_q_p"] = (
              (log_q_z_given_x - log_p_z).detach().cpu()
          )
  
      elif self.model_type == "nll":
          loss = -log_p_x_given_z
  
      if return_info:
          return loss, info_dict
      else:
          return loss
  
imports: from typing import Dict, Optional
  import numpy
  import torch
  from pytorch_lightning import Trainer
  from nemo.collections.common.losses import NLLLoss
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_config import MTBottleneckModelConfig
  from nemo.collections.nlp.models.machine_translation.mt_enc_dec_model import MTEncDecModel
  from nemo.core.classes.common import typecheck
  from nemo.utils import timers
label_desc: Use image node_id 6 for calling the MTBottleneckModel obj's underlying member method code with example usage: obj.loss(z, z_mean, z_logv, z_mask, tgt_log_probs, tgt, tgt_mask, tgt_labels, train, return_info) and returns: loss,  info_dict, loss