text: import torch
  import torch.nn.functional as F
  from torch.nn.modules.loss import _Loss
  from torch.distributions import MultivariateNormal as MVN
  
  
  class ReweightL2(_Loss):
      def __init__(self, train_dist, reweight='inverse'):
          super(ReweightL2, self).__init__()
          self.reweight = reweight
          self.train_dist = train_dist
  
      def forward(self, pred, target):
          reweight = self.reweight
          prob = self.train_dist.log_prob(target).exp().squeeze(-1)
          if reweight == 'inverse':
              inv_prob = prob.pow(-1)
          elif reweight == 'sqrt_inv':
              inv_prob = prob.pow(-0.5)
          else:
              raise NotImplementedError
          inv_prob = inv_prob / inv_prob.sum()
          loss = F.mse_loss(pred, target, reduction='none').sum(-1) * inv_prob
          loss = loss.sum()
          return loss
  
  
  class GAILossMD(_Loss):
      """
      Multi-Dimension version GAI, compatible with 1-D GAI
      """
  
      def __init__(self, init_noise_sigma, gmm):
          super(GAILossMD, self).__init__()
          self.gmm = gmm
          self.gmm = {k: torch.tensor(self.gmm[k]) for k in self.gmm}
          self.noise_sigma = torch.nn.Parameter(torch.tensor(init_noise_sigma))
  
      def forward(self, pred, target):
          noise_var = self.noise_sigma ** 2
          loss = gai_loss_md(pred, target, self.gmm, noise_var)
          return loss
  
  
  def gai_loss_md(pred, target, gmm, noise_var):
      I = torch.eye(pred.shape[-1])
      mse_term = -MVN(pred, noise_var*I).log_prob(target)
      balancing_term = MVN(gmm['means'], gmm['variances']+noise_var*I).log_prob(pred.unsqueeze(1)) + gmm['weights'].log()
      balancing_term = torch.logsumexp(balancing_term, dim=1)
      loss = mse_term + balancing_term
      loss = loss * (2 * noise_var).detach()
      return loss.mean()
  
  
  class BMCLossMD(_Loss):
      """
      Multi-Dimension version BMC, compatible with 1-D BMC
      """
  
      def __init__(self, init_noise_sigma):
          super(BMCLossMD, self).__init__()
          self.noise_sigma = torch.nn.Parameter(torch.tensor(init_noise_sigma))
  
      def forward(self, pred, target):
          noise_var = self.noise_sigma ** 2
          loss = bmc_loss_md(pred, target, noise_var)
          return loss
  
  
  def bmc_loss_md(pred, target, noise_var):
      I = torch.eye(pred.shape[-1])
      logits = MVN(pred.unsqueeze(1), noise_var*I).log_prob(target.unsqueeze(0))
      loss = F.cross_entropy(logits, torch.arange(pred.shape[0]))
      loss = loss * (2 * noise_var).detach()
      return loss
  