code: def infoNCE_local_intra_node(
      z1_n, z2_n, tau=0.5, norm=True, batch=None
  ):
      """
      Args:
          z1_n: Tensor of shape [n_nodes, z_dim].
          z2_n: Tensor of shape [n_nodes, z_dim].
          tau: Float. Usually in (0,1].
          norm: Boolean. Whether to apply normlization.
          batch: Tensor of shape [batch_size]
      """
  
      def sim(z1: torch.Tensor, z2: torch.Tensor):
          if norm:
              z1 = F.normalize(z1)
              z2 = F.normalize(z2)
          return torch.mm(z1, z2.t())
  
      exp = lambda x: torch.exp(x / tau)
      if batch is not None:
          batch_size = batch.size(0)
          num_nodes = z1_n.size(0)
          indices = torch.arange(0, num_nodes).to(z1_n.device)
          losses = []
          for i in range(0, num_nodes, batch_size):
              mask = indices[i : i + batch_size]
              refl_sim = exp(sim(z1_n[mask], z1_n))
              between_sim = exp(sim(z1_n[mask], z2_n))
              losses.append(
                  -torch.log(
                      between_sim[:, i : i + batch_size].diag()
                      / (
                          refl_sim.sum(1)
                          + between_sim.sum(1)
                          - refl_sim[:, i : i + batch_size].diag()
                      )
                  )
              )
          losses = torch.cat(losses)
          return losses.mean()
  
      refl_sim = exp(sim(z1_n, z1_n))
      between_sim = exp(sim(z1_n, z2_n))
  
      pos_sim = between_sim.diag()
      intra_sim = refl_sim.sum(1) - refl_sim.diag()
      inter_pos_sim = between_sim.sum(1)
  
      loss = pos_sim / (intra_sim + inter_pos_sim)
      loss = -torch.log(loss).mean()
  
      return loss
  
imports: import itertools
  import torch
  import torch.nn.functional
label_desc: Use image node_id 2 for calling a global function with example usage: infoNCE_local_intra_node(z1_n, z2_n, tau, norm, batch) and returns: loss, torch, losses