code: def NCE_loss(zs=None, zs_n=None, batch=None, sigma=None, **kwargs):
      """The InfoNCE (NT-XENT) loss in contrastive learning.
  
      Args:
          zs (list, optipnal): List of tensors of shape [batch_size, z_dim].
          zs_n (list, optional): List of tensors of shape [nodes, z_dim].
          batch (Tensor, optional): Required when both :obj:`zs` and :obj:`zs_n` are given.
          sigma (ndarray, optional): A 2D-array of shape [:obj:`n_views`, :obj:`n_views`] with boolean
              values, indicating contrast between which two views are computed. Only required
              when number of views is greater than 2. If :obj:`sigma[i][j]` = :obj:`True`,
              infoNCE between :math:`view_i` and :math:`view_j` will be computed.
          tau (int, optional): The temperature used in NT-XENT.
  
      :rtype: :class:`Tensor`
      """
      assert zs is not None or zs_n is not None
  
      if "tau" in kwargs:
          tau = kwargs["tau"]
      else:
          tau = 0.5
  
      if "norm" in kwargs:
          norm = kwargs["norm"]
      else:
          norm = True
  
      mean = kwargs["mean"] if "mean" in kwargs else True
  
      if zs_n is not None:
          if zs is None:
              # InfoNCE in GRACE
              assert len(zs_n) == 2
              return (
                  infoNCE_local_intra_node(
                      zs_n[0], zs_n[1], tau, norm, batch
                  )
                  + infoNCE_local_intra_node(
                      zs_n[1], zs_n[0], tau, norm, batch
                  )
              ) * 0.5
          else:
              assert len(zs_n) == len(zs)
              assert batch is not None
  
              if len(zs) == 1:
                  return infoNCE_local_global(
                      zs[0], zs_n[0], batch, tau, norm
                  )
              elif len(zs) == 2:
                  return infoNCE_local_global(
                      zs[0], zs_n[1], batch, tau, norm
                  ) + infoNCE_local_global(
                      zs[1], zs_n[0], batch, tau, norm
                  )
              else:
                  assert len(zs) == len(sigma)
                  loss = 0
                  for i, j in itertools.combinations(range(len(zs)), 2):
                      if sigma[i][j]:
                          loss += infoNCE_local_global(
                              zs[i], zs_n[j], batch, tau, norm
                          ) + infoNCE_local_global(
                              zs[j], zs_n[i], batch, tau, norm
                          )
                  return loss
  
      if len(zs) == 2:
          return NT_Xent(zs[0], zs[1], tau, norm)
      elif len(zs) > 2:
          assert len(zs) == len(sigma)
          loss = 0
          for i, j in itertools.combinations(range(len(zs)), 2):
              if sigma[i][j]:
                  loss += NT_Xent(zs[i], zs[j], tau, norm)
          return loss
  
imports: import itertools
  import torch
  import torch.nn.functional
label_desc: Use image node_id 1 for calling a global function with example usage: NCE_loss(zs, zs_n, batch, sigma) and returns: NT_Xent, unknown, loss, infoNCE_local_global, unknown, loss