text: # Copyright 2021 solo-learn development team.
  
  # Permission is hereby granted, free of charge, to any person obtaining a copy of
  # this software and associated documentation files (the "Software"), to deal in
  # the Software without restriction, including without limitation the rights to use,
  # copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the
  # Software, and to permit persons to whom the Software is furnished to do so,
  # subject to the following conditions:
  
  # The above copyright notice and this permission notice shall be included in all copies
  # or substantial portions of the Software.
  
  # THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED,
  # INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR
  # PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE
  # FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR
  # OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER
  # DEALINGS IN THE SOFTWARE.
  
  import argparse
  from typing import Any, Dict, List, Sequence, Tuple
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from solo.losses.moco import moco_loss_func
  # from solo.losses.simclr import simclr_loss_func
  # from solo.methods.base import  BaseDistillationMethod
  from solo.methods.base_for_adversarial_training import BaseDistillationATMethod
  from solo.utils.momentum import initialize_momentum_params
  from solo.utils.misc import gather
  
  from torchvision import models
  from solo.utils.metrics import accuracy_at_k, weighted_mean
  
  
  class MoCoV2KDAT(BaseDistillationATMethod):
      queue: torch.Tensor
  
      def __init__(
          self,
          proj_output_dim: int,
          proj_hidden_dim: int,
          temperature: float,
          queue_size: int,
          teacher_logit_fix: bool,
          loss_type: str,
          projector_ablation: str,
          epsilon: int = 8,
          num_steps: int = 5,
          step_size: int = 2,
          trades_k: float = 3, 
          aux_data: bool = False,
          augmentation_ablation: bool = False,
          expriment_code: str = "000",
          **kwargs
      ):
          """Implements MoCo V2+ (https://arxiv.org/abs/2011.10566).
  
          Args:
              proj_output_dim (int): number of dimensions of projected features.
              proj_hidden_dim (int): number of neurons of the hidden layers of the projector.
              temperature (float): temperature for the softmax in the contrastive loss.
              queue_size (int): number of samples to keep in the queue.
          """
  
          super().__init__(**kwargs)
  
          self.temperature = temperature
          self.queue_size = queue_size
          self.loss_type = loss_type
          self.projector_ablation = projector_ablation
          self.trades_k = trades_k
          self.aux_data = aux_data
  
          self.augmentation_ablation = augmentation_ablation
          self.expriment_code = expriment_code
  
  
          self.epsilon = epsilon/255.
          self.num_steps = num_steps
          self.step_size = step_size/255.
  
          # projector
          self.projector = nn.Sequential(
              nn.Linear(self.features_dim, proj_hidden_dim),
              nn.ReLU(),
              nn.Linear(proj_hidden_dim, proj_output_dim),
          )
  
          
          # momentum projector
          self.momentum_projector = nn.Sequential(
              nn.Linear(self.features_dim, proj_hidden_dim),
              nn.ReLU(),
              nn.Linear(proj_hidden_dim, proj_output_dim),
          )
  
          # create the queue
          self.register_buffer("queue", torch.randn(2, self.features_dim, queue_size))
          self.queue = nn.functional.normalize(self.queue, dim=1)
          self.register_buffer("queue_ptr", torch.zeros(1, dtype=torch.long))
  
          # To remove both projector and projector pretraining weight loading
          if self.projector_ablation == "remove":
              self.projector = nn.Identity()
              self.momentum_projector = nn.Identity()
          
          elif self.projector_ablation == "correspond":
              # Make StudentNetwork have same weight with Teacher
              self.projector.load_state_dict(self.projector_state_dict)
              initialize_momentum_params(self.projector, self.momentum_projector)
              self.momentum_projector.load_state_dict(self.projector_state_dict)
          
          elif self.projector_ablation == "only_student":
              self.projector.load_state_dict(self.projector_state_dict)
              initialize_momentum_params(self.projector, self.momentum_projector)
              self.momentum_projector = nn.Identity()
          
          elif self.projector_ablation == "only_teacher":
              self.projector = nn.Identity()
              initialize_momentum_params(self.projector, self.momentum_projector)
              self.momentum_projector.load_state_dict(self.projector_state_dict)
  
  
          self.teacher_logit_fix = teacher_logit_fix
          if teacher_logit_fix:
          
              teacher_logit = torch.load(f"{self.teacher_ckpt_dir}.features_collection.pt")
  
              self.train_features = teacher_logit["train_features"]
              self.train_features_label = teacher_logit["train_features_label"]
  
              idx_unshuffle = torch.argsort(teacher_logit["features_idx"])
              self.train_features = self.train_features[idx_unshuffle]
              self.train_features_label = self.train_features_label[idx_unshuffle]
  
  
  
      @staticmethod
      def add_model_specific_args(parent_parser: argparse.ArgumentParser) -> argparse.ArgumentParser:
          parent_parser = super(MoCoV2KDAT, MoCoV2KDAT).add_model_specific_args(parent_parser)
          parser = parent_parser.add_argument_group("mocov2_kd_at")
  
          # projector
          parser.add_argument("--proj_output_dim", type=int, default=128)
          parser.add_argument("--proj_hidden_dim", type=int, default=2048)
  
          # parameters
          parser.add_argument("--temperature", type=float, default=0.1)
  
          # queue settings
          parser.add_argument("--queue_size", default=65536, type=int)
  
  
          parser.add_argument("--teacher_logit_fix", action="store_true")
  
          # choose loss to train the model
          SUPPORT_LOSS = ["ce2gt_ae2ce",
                          "ae2gt_ae2ce",
                          "ce2gt_ae2ce_infonce",
                          "pgd",
                          "ae2gt_infonce",
  
                          "ae2gt_kl",
                          "ce2gt_ae2ce_kl",
                          "cs.ce2gt._kl.ae2ce.",
                          "kl.ce2gt._cs.ae2ce.",
                          "cs.ce2gt._infonce.ae2ce.",
                          "infonce.ce2gt._cs.ae2ce.",
                          "infonce.ce2gt._kl.ae2ce.",
                          "kl.ce2gt._infonce.ae2ce."
                          ]
  
          parser.add_argument("--loss_type", default="ce2gt_ae2ce", type=str, choices=SUPPORT_LOSS)
  
  
          # projector exploration
          PROJECTOR_ABLATION = ["remove", "same", "correspond", 'only_student', 'only_teacher']
          parser.add_argument("--projector_ablation", default="remove", type=str, choices=PROJECTOR_ABLATION)
  
  
          # training adversarial hyper parameter
          parser.add_argument("--epsilon", type=int, default=8)
          parser.add_argument("--step_size", type=int, default=2)
          parser.add_argument("--num_steps", type=int, default=5)
  
          # for loss factor
          parser.add_argument("--trades_k", type=float, default=2)
  
          # for augmentation ablation study
          parser.add_argument("--augmentation_ablation", action="store_true")
          parser.add_argument("--expriment_code", default="00", type=str)
          
  
  
  
  
          return parent_parser
  
      @property
      def learnable_params(self) -> List[dict]:
          """Adds projector parameters together with parent's learnable parameters.
  
          Returns:
              List[dict]: list of learnable parameters.
          """
  
          extra_learnable_params = [{"params": self.projector.parameters()}]
          return super().learnable_params + extra_learnable_params
  
      @property
      def momentum_pairs(self) -> List[Tuple[Any, Any]]:
          """Adds (projector, momentum_projector) to the parent's momentum pairs.
  
          Returns:
              List[Tuple[Any, Any]]: list of momentum pairs.
          """
  
          extra_momentum_pairs = [(self.projector, self.momentum_projector)]
          return super().momentum_pairs + extra_momentum_pairs
  
      @torch.no_grad()
      def _dequeue_and_enqueue(self, keys: torch.Tensor):
          """Adds new samples and removes old samples from the queue in a fifo manner.
  
          Args:
              keys (torch.Tensor): output features of the momentum backbone.
          """
  
          batch_size = keys.shape[1]
          ptr = int(self.queue_ptr)  # type: ignore
          assert self.queue_size % batch_size == 0  # for simplicity
  
          # replace the keys at ptr (dequeue and enqueue)
          keys = keys.permute(0, 2, 1)
          self.queue[:, :, ptr : ptr + batch_size] = keys
          ptr = (ptr + batch_size) % self.queue_size  # move pointer
          self.queue_ptr[0] = ptr  # type: ignore
  
      # def forward(self, X: torch.Tensor, *args, **kwargs) -> Dict[str, Any]:
      #     """Performs the forward pass of the online backbone and projector.
  
      #     Args:
      #         X (torch.Tensor): a batch of images in the tensor format.
  
      #     Returns:
      #         Dict[str, Any]: a dict containing the outputs of the parent and the projected features.
      #     """
  
      #     out = super().forward(X, *args, **kwargs)
      #     z = F.normalize(self.projector(out["feats"]), dim=-1)
      #     return {**out, "z": z}
  
      def forward(self, X: torch.Tensor):
          """Performs the forward pass of the online backbone and projector.
  
          Args:
              X (torch.Tensor): a batch of images in the tensor format.
  
          Returns:
              Dict[str, Any]: a dict containing the outputs of the parent and the projected features.
          """
          out = self.backbone(X)
  
          return out
  
      def training_step(self, batch: Sequence[Any], batch_idx: int) -> torch.Tensor:
          """
          Training step for MoCo reusing BaseMomentumMethod training step.
  
          Args:
              batch (Sequence[Any]): a batch of data in the
                  format of [img_indexes, [X], Y], where [X] is a list of size self.num_large_crops
                  containing batches of images.
              batch_idx (int): index of the batch.
  
          Returns:
              torch.Tensor: total loss composed of MOCO loss and classification loss.
  
          """
  
  
          self.momentum_backbone.eval()
  
  
          # import ipdb; ipdb.set_trace()
          if self.aux_data:
              image_tau1, image_weak = batch[0]
              targets = batch[1]
  
          else:
              image_tau1, image_weak = batch[1]
              targets = batch[2]
  
  
          ############################################################################
          # Adversarial Training Trades-like loss (CAT)
          ############################################################################
          if self.augmentation_ablation:
              experiment_code = [int(code) for code in self.expriment_code]
  
              with torch.no_grad():
                  teacher_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
              if experiment_code[0] == 0:
                  AE_generation_image = image_weak
              elif experiment_code[0] == 1:
                  AE_generation_image = image_tau1
  
  
              # AE_generation_image = image_weak
              image_AE = self.generate_training_AE(AE_generation_image, teacher_target)
  
              if experiment_code[1] == 0:
                  Student_Clean = image_weak
              elif experiment_code[1] == 1:
                  Student_Clean = image_tau1
  
              image_CAT = torch.cat([Student_Clean, image_AE])
              logits_all = self.projector(self.backbone(image_CAT))
              bs = image_weak.size(0)
              student_logtis_clean = logits_all[:bs]
              student_logits_AE = logits_all[bs:]
  
              # Cosine Similarity loss
              adv_loss = -F.cosine_similarity(student_logtis_clean, teacher_target).mean()
              adv_loss += -self.trades_k*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
  
  
          ############################################################################
          # Adversarial Training Trades-like loss (CAT)
          ############################################################################
  
          else:
          ############################################################################
          # Adversarial Training (CAT)
          ############################################################################
          # import ipdb; ipdb.set_trace()
              if self.teacher_logit_fix:
  
                  away_target = self.train_features[batch[0]].cuda()
              else:
                  away_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
              AE_generation_image = image_weak
              image_AE = self.generate_training_AE(AE_generation_image, away_target)
  
              image_CAT = torch.cat([image_weak, image_AE])
              logits_all = self.projector(self.backbone(image_CAT))
              bs = image_weak.size(0)
              student_logtis_clean = logits_all[:bs]
              student_logits_AE = logits_all[bs:]
  
              # Cosine Similarity loss
              # adv_loss = -F.cosine_similarity(student_logits_AE, away_target).mean()
              if self.loss_type == "ce2gt_ae2ce":
                  adv_loss = -F.cosine_similarity(student_logtis_clean, away_target).mean()
                  adv_loss += -self.trades_k*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
              elif self.loss_type == "ae2gt_ae2ce":
                  adv_loss = -F.cosine_similarity(student_logits_AE, away_target).mean()
                  adv_loss += -self.trades_k*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
              elif self.loss_type == "pgd":
                  adv_loss = -F.cosine_similarity(student_logits_AE, away_target).mean()
              elif self.loss_type == "ce2gt_ae2ce_infonce":
  
                  adv_loss = simclr_loss_func(
                      student_logtis_clean, away_target,
                      temperature=self.temperature,
                  )
  
                  adv_loss += simclr_loss_func(
                      student_logits_AE, student_logtis_clean,
                      temperature=self.temperature,
                  )
              elif self.loss_type == "ae2gt_infonce":
  
                  adv_loss = simclr_loss_func(
                      student_logits_AE, away_target,
                      temperature=self.temperature,
                  )
  
              elif self.loss_type == "ae2gt_kl":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = criterion_kl(F.log_softmax(student_logits_AE, dim=1),
                                         F.softmax(away_target, dim=1))
  
              elif self.loss_type == "ce2gt_ae2ce_kl":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = criterion_kl(F.log_softmax(student_logtis_clean, dim=1),
                                         F.softmax(away_target, dim=1)) +\
                              self.trades_k*criterion_kl(F.log_softmax(student_logits_AE, dim=1),
                                         F.softmax(away_target, dim=1))
  
              elif self.loss_type == "cs.ce2gt._kl.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = -F.cosine_similarity(student_logtis_clean, away_target).mean() +\
                              self.trades_k*criterion_kl(F.log_softmax(student_logits_AE, dim=1),
                                         F.softmax(away_target, dim=1))
  
              elif self.loss_type == "kl.ce2gt._cs.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = criterion_kl(F.log_softmax(student_logtis_clean, dim=1),
                                         F.softmax(away_target, dim=1)) +\
                              -self.trades_k*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
  
  
              elif self.loss_type == "cs.ce2gt._infonce.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = -F.cosine_similarity(student_logtis_clean, away_target).mean() +\
                              self.trades_k*simclr_loss_func(
                      student_logits_AE, away_target,
                      temperature=self.temperature,
                  )
  
              elif self.loss_type == "infonce.ce2gt._cs.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = simclr_loss_func(
                      student_logtis_clean, away_target,
                      temperature=self.temperature,
                  ) +\
                              -self.trades_k*F.cosine_similarity(student_logits_AE, student_logtis_clean).mean()
  
  
              elif self.loss_type == "infonce.ce2gt._kl.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = simclr_loss_func(
                      student_logtis_clean, away_target,
                      temperature=self.temperature,
                  )  +\
                              self.trades_k*criterion_kl(F.log_softmax(student_logits_AE, dim=1),
                                         F.softmax(away_target, dim=1))
  
              elif self.loss_type == "kl.ce2gt._infonce.ae2ce.":
                  criterion_kl = nn.KLDivLoss(size_average=False)
  
                  adv_loss = criterion_kl(F.log_softmax(student_logtis_clean, dim=1),
                                         F.softmax(away_target, dim=1)) +\
                              -self.trades_k*simclr_loss_func(
                      student_logits_AE, away_target,
                      temperature=self.temperature,
                  )
  
          ############################################################################
          # Adversarial Training (CAT)
          ############################################################################
  
  
  
  
  
  
          ############################################################################
          # Adversarial Training (CODE Experiment)
          ############################################################################
  
          # # NOTE: Check Wether to Make Student Start Same Weight with Teacher
          # # CODE Definition. 0: Weak augmentation 1: Tau1 augmentation
          # # CS_Teacher$1_Student-AE_CS_Teacher$2_Student$3, so case 1 is 110
          # # 1. CS_TeacherTau1_Student-AE_CS_TeacherTau1_StudentWeak
          # # 110
          # # with torch.no_grad():
          # #     logits_TeacherTau1 = self.projector(self.momentum_backbone(image_tau1))
  
          # # image_AE = self.generate_training_AE(image_weak, logits_TeacherTau1)
          # # student_logit = self.projector(self.backbone(image_AE))
          # # adv_loss = -F.cosine_similarity(student_logit, logits_TeacherTau1).mean()
  
          # # 2. CS_TeacherTau1_Student-AE_CS_TeacherTau1_StudentTau1
          # # 111
          # # with torch.no_grad():
          # #     logits_TeacherTau1 = self.projector(self.momentum_backbone(image_tau1))
  
          # # image_AE = self.generate_training_AE(image_tau1, logits_TeacherTau1)
          # # student_logit = self.projector(self.backbone(image_AE))
          # # adv_loss = -F.cosine_similarity(student_logit, logits_TeacherTau1).mean()
  
          # # CODE Based experiment
          # # 110
          # # 111
          # # experiment_code = [1, 0, 1]
          # # experiment_code = [1, 0, 0]
  
          # experiment_code = [0, 0, 0]
          # # experiment_code = [0, 1, 1]
          # # experiment_code = [0, 0, 1]
          # # experiment_code = [0, 1, 0]
  
          # with torch.no_grad():
          #     # away target
          #     if experiment_code[1] == 0:
          #         # away_target = self.projector(self.momentum_backbone(image_weak))
          #         away_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
          #     elif experiment_code[1] == 1:
          #         # away_target = self.projector(self.momentum_backbone(image_tau1))
          #         away_target = self.momentum_projector(self.momentum_backbone(image_tau1))
  
  
          #     # Loss learning target
          #     if experiment_code[0] == experiment_code[1]:
          #         learning_target = away_target
          #     elif experiment_code[0] == 0:
          #         # learning_target = self.projector(self.momentum_backbone(image_weak))
          #         learning_target = self.momentum_projector(self.momentum_backbone(image_weak))
  
          #     elif experiment_code[0] == 1:
          #         # learning_target = self.projector(self.momentum_backbone(image_tau1))
          #         learning_target = self.momentum_projector(self.momentum_backbone(image_tau1))
  
  
          # if experiment_code[2] == 0:
          #     AE_generation_image = image_weak
          # elif experiment_code[2] == 1:
          #     AE_generation_image = image_tau1
  
          # image_AE = self.generate_training_AE(AE_generation_image, away_target)
          # student_logit = self.projector(self.backbone(image_AE))
          # adv_loss = -F.cosine_similarity(student_logit, learning_target).mean()
  
          ############################################################################
          # Adversarial Training (CODE Experiment)
          ############################################################################
  
  
  
  
          ############################################################################
          # Online clean classifier training
          ############################################################################
          # Bug Fix: train classifier using evaluation mode
          self.backbone.eval()
          # import ipdb; ipdb.set_trace()
          outs_image_weak = self._base_shared_step(image_weak, targets)
          self.backbone.train()
          metrics = {
              "train_class_loss": outs_image_weak["loss"],
              "train_acc1": outs_image_weak["acc1"],
              "train_acc5": outs_image_weak["acc5"],
          }
          class_loss_clean = outs_image_weak["loss"]
          self.log_dict(metrics, on_epoch=True)
          ############################################################################
          # Online clean classifier training
          ############################################################################
  
  
  
  
          ############################################################################
          # Online adv classifier training
          ############################################################################
          # # Bug Fix: train classifier using evaluation mode
          self.backbone.eval()
          AE_feats = self.backbone(image_AE)
          logits = self.classifier_adv(AE_feats.detach())
          self.backbone.train()
  
          class_loss_adv = F.cross_entropy(logits, targets, ignore_index=-1)
          # handle when the number of classes is smaller than 5
          top_k_max = min(5, logits.size(1))
          acc1, acc5 = accuracy_at_k(logits, targets, top_k=(1, top_k_max))
          metrics = {
              "train_class_loss_adv_classifier": class_loss_adv,
              "train_acc1_adv_classifier": acc1,
              "train_acc5_adv_classifier": acc5,
          }
          class_loss_clean = outs_image_weak["loss"]
          self.log_dict(metrics, on_epoch=True)
          ############################################################################
          # Online adv classifier training
          ############################################################################
  
  
  
          self.log("adv_loss", adv_loss, on_epoch=True, sync_dist=True)
  
          # return adv_loss + class_loss_adv + class_loss_clean
          return adv_loss + class_loss_clean
  
  
  
      def generate_training_AE(self, image: torch.Tensor, away_target: torch.Tensor):
          """
          images_org: weak aug
          away_target: from teacher
          """
  
          # self.epsilon = 8/255.
          # self.num_steps = 5
          # self.step_size = 2/255.
  
          x_cl = image.clone().detach()
  
          # if self.rand:
          x_cl = x_cl + torch.zeros_like(image).uniform_(-self.epsilon, self.epsilon)
  
          # f_ori_proj = self.model(images_org).detach()
          # Change the attack process of model to eval
          self.backbone.eval()
  
          for i in range(self.num_steps):
              x_cl.requires_grad_()
              with torch.enable_grad():
                  f_proj = self.projector(self.backbone(x_cl))
  
                  # loss_contrast = -F.cosine_similarity(f_proj, f_ori_proj, dim=1).mean()
                  loss_contrast = -F.cosine_similarity(f_proj, away_target, dim=1).sum() *256
                  loss = loss_contrast
  
              # import ipdb ;ipdb.set_trace()
              grad_x_cl = torch.autograd.grad(loss, x_cl)[0]
              # grad_x_cl = torch.autograd.grad(loss, x_cl, grad_outputs=torch.ones_like(loss))[0]
              x_cl = x_cl.detach() + self.step_size * torch.sign(grad_x_cl.detach())
  
              # remove the clamp in for the image comparision
              x_cl = torch.min(torch.max(x_cl, image - self.epsilon), image + self.epsilon)
              x_cl = torch.clamp(x_cl, 0, 1)
  
          self.backbone.train()
  
          return x_cl
  
  
  
  
  
  
  
  
  
  def simclr_loss_func(
      z1: torch.Tensor,
      z2: torch.Tensor,
      temperature: float = 0.1,
      extra_pos_mask=None,
  ) -> torch.Tensor:
      """Computes SimCLR's loss given batch of projected features z1 from view 1 and
      projected features z2 from view 2.
      Args:
          z1 (torch.Tensor): NxD Tensor containing projected features from view 1.
          z2 (torch.Tensor): NxD Tensor containing projected features from view 2.
          temperature (float): temperature factor for the loss. Defaults to 0.1.
          extra_pos_mask (Optional[torch.Tensor]): boolean mask containing extra positives other
              than normal across-view positives. Defaults to None.
      Returns:
          torch.Tensor: SimCLR loss.
      """
  
      device = z1.device
  
      b = z1.size(0)
      z = torch.cat((z1, z2), dim=0)
      z = F.normalize(z, dim=-1)
  
      logits = torch.einsum("if, jf -> ij", z, z) / temperature
      logits_max, _ = torch.max(logits, dim=1, keepdim=True)
      logits = logits - logits_max.detach()
  
      # positive mask are matches i, j (i from aug1, j from aug2), where i == j and matches j, i
      pos_mask = torch.zeros((2 * b, 2 * b), dtype=torch.bool, device=device)
      pos_mask[:, b:].fill_diagonal_(True)
      pos_mask[b:, :].fill_diagonal_(True)
  
      # if we have extra "positives"
      if extra_pos_mask is not None:
          pos_mask = torch.bitwise_or(pos_mask, extra_pos_mask)
  
      # all matches excluding the main diagonal
      logit_mask = torch.ones_like(pos_mask, device=device).fill_diagonal_(0)
  
      exp_logits = torch.exp(logits) * logit_mask
      log_prob = logits - torch.log(exp_logits.sum(1, keepdim=True))
  
      # compute mean of log-likelihood over positives
      mean_log_prob_pos = (pos_mask * log_prob).sum(1) / pos_mask.sum(1)
      # loss
      loss = -mean_log_prob_pos.mean()
      return loss
  
  