class_docstr: Based on <https://arxiv.org/abs/2003.02645>
  Implements A-MIM loss with a unit Normal anchor.
  A-MIM - asymmetric MIM (without sampling)
class_docstr_tok: ["Based","on","<","https",":","\/\/arxiv.org\/abs\/2003.02645",">","Implements","A-MIM","loss","with","a","unit","Normal","anchor",".","A-MIM","-","asymmetric","MIM","(","without","sampling",")"]
docstr: We expect input shapes to be [S x B x H] for Sequence, Batch, Hidden sizes (due to tensor parallel support).
  We return a dictionary with dimensions [B x S x H], [B x S], [B], or [].
  
  Implement your own loss calculations. Must return "loss" key.
  loss shape - [B x S] for Batch, Sequence sizes
  batch_data - a dictionary of additional data that can be used to calculate loss
docstr_tok: ["We","expect","input","shapes","to","be","[","S","x","B","x","H","]","for","Sequence",",","Batch",",","Hidden","sizes","(","due","to","tensor","parallel","support",")",".","We","return","a","dictionary","with","dimensions","[","B","x","S","x","H","]",",","[","B","x","S","]",",","[","B","]",",","or","[","]",".","Implement","your","own","loss","calculations",".","Must","return","``","loss","''","key",".","loss","shape","-","[","B","x","S","]","for","Batch",",","Sequence","sizes","batch_data","-","a","dictionary","of","additional","data","that","can","be","used","to","calculate","loss"]
code: def _loss(self, inputs, batch_data=None):
      """
      We expect input shapes to be [S x B x H] for Sequence, Batch, Hidden sizes (due to tensor parallel support).
      We return a dictionary with dimensions [B x S x H], [B x S], [B], or [].
  
      Implement your own loss calculations. Must return "loss" key.
      loss shape - [B x S] for Batch, Sequence sizes
      batch_data - a dictionary of additional data that can be used to calculate loss
      """
      z = inputs["z"]
      # get posterior
      log_prob_q_z_given_x = inputs["z_log_prob"]
      # compute log prob of anchor a unit Normal distribution
      log_prob_P_z = -0.5 * (math.log(2 * math.pi) + z.pow(2))
      # aggregate over hidden dimension, default is sum
      log_prob_P_z = getattr(
          log_prob_P_z, self.hidden_aggregation_method
      )(dim=-1)
  
      # A-MIM loss = log_p_x_given_z - 0.5 * (log_prob_P_z + log_prob_q_z_given_x)
      # here we return only the hidden loss part
      loss = -0.5 * (log_prob_P_z + log_prob_q_z_given_x)
  
      # return losses shaped [B x S]
      return {
          "loss": loss.transpose(0, 1),
          "log_prob_P_z": log_prob_P_z.transpose(0, 1),
          "log_prob_q_z_given_x": log_prob_q_z_given_x.transpose(0, 1),
      }
  
imports: import math
  import torch
label_desc: Use image node_id 3 for calling the MegatronAMIMHiddenLoss obj's underlying member method code with example usage: obj._loss(inputs, batch_data) and returns: dict