class_docstr: Boosted version of MultipleNegativesRankingLoss (https://arxiv.org/pdf/1705.00652.pdf) by GradCache (https://arxiv.org/pdf/2101.06983.pdf).
  
  Constrastive learning (here our MNRL loss) with in-batch negatives is usually hard to work with large batch sizes due to (GPU) memory limitation.
  Even with batch-scaling methods like gradient-scaling, it cannot work either. This is because the in-batch negatives make the data points within
  the same batch non-independent and thus the batch cannot be broke down into mini-batches. GradCache is a smart way to solve this problem.
  It achieves the goal by dividing the computation into two stages of embedding and loss calculation, which both can be scaled by mini-batches.
  As a result, memory of constant size (e.g. that works with batch size = 32) can now process much larger batches (e.g. 65536).
  
  In detail:
      (1) It first does a quick embedding step without gradients/computation graphs to get all the embeddings;
      (2) Calculate the loss, backward up to the embeddings and cache the gradients wrt. to the embeddings;
      (3) A 2nd embedding step with gradients/computation graphs and connect the cached gradients into the backward chain.
  
  Notes: All steps are done with mini-batches. In the original implementation of GradCache, (2) is not done in mini-batches and
  requires a lot memory when batch size large. One drawback is about the speed. GradCache will sacrifice around 20% computation time according to the paper.
  
  Example:
  
      from sentence_transformers import SentenceTransformer
      train_dataloader = DataLoader(train_examples, shuffle=True, batch_size=1024)  # Here we can try much larger batch sizes!
      train_loss = losses.CachedMultipleNegativesRankingLoss(model=model, mini_batch_size: int = 32)
docstr: Calculate the cross-entropy loss and cache the gradients wrt. the embeddings.
docstr_tok: ["Calculate","the","cross-entropy","loss","and","cache","the","gradients","wrt",".","the","embeddings","."]
code: def calculate_loss_and_cache_gradients(
      self, reps: List[List[Tensor]]
  ) -> Tensor:
      """Calculate the cross-entropy loss and cache the gradients wrt. the embeddings."""
      embeddings_a = torch.cat(reps[0])  # (bsz, hdim)
      embeddings_b = torch.cat(
          [torch.cat(r) for r in reps[1:]]
      )  # ((1 + nneg) * bsz, hdim)
  
      batch_size = len(embeddings_a)
      labels = torch.tensor(
          range(batch_size),
          dtype=torch.long,
          device=embeddings_a.device,
      )  # (bsz, (1 + nneg) * bsz)  Example a[i] should match with b[i]
      losses: List[torch.Tensor] = []
      for b in tqdm.trange(
          0,
          batch_size,
          self.mini_batch_size,
          desc="Preparing caches",
          disable=not self.show_progress_bar,
      ):
          e = b + self.mini_batch_size
          scores: Tensor = (
              self.similarity_fct(embeddings_a[b:e], embeddings_b)
              * self.scale
          )
          loss_mbatch: torch.Tensor = (
              self.cross_entropy_loss(scores, labels[b:e])
              * len(scores)
              / batch_size
          )
          loss_mbatch.backward()
          losses.append(loss_mbatch.detach())
  
      loss = sum(losses).requires_grad_()
  
      self.cache = [
          [r.grad for r in rs] for rs in reps
      ]  # e.g. 3 * bsz/mbsz * (mbsz, hdim)
  
      return loss
  
imports: from __future__ import annotations
  from contextlib import nullcontext
  from functools import partial
  import torch
  from torch import nn, Tensor
  from torch.utils.checkpoint import get_device_states, set_device_states
  from typing import Iterable, Dict, Iterator, List, Optional, Tuple
  from sentence_transformers import SentenceTransformer
  from sentence_transformers import util
  import tqdm
label_desc: Use image node_id 4 for calling the CachedMultipleNegativesRankingLoss obj's underlying member method code with example usage: obj.calculate_loss_and_cache_gradients(reps) and returns: loss