code: def JSE_loss(
      zs, zs_n=None, batch=None, sigma=None, neg_by_crpt=False, **kwargs
  ):
      """The Jensen-Shannon Estimator of Mutual Information used in contrastive learning. The
      implementation follows the paper `Learning deep representations by mutual information
      estimation and maximization <https://arxiv.org/abs/1808.06670>`_.
  
      .. note::
          The JSE loss implementation can produce negative values because a :obj:`-2log2` shift is
          added to the computation of JSE, for the sake of consistency with other f-convergence
          losses.
  
      Args:
          zs (list, optional): 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 [n_views, 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`,
              JSE between :math:`view_i` and :math:`view_j` will be computed.
          neg_by_crpt (bool, optional): The mode to obtain negative samples in JSE. If True,
              obtain negative samples by performing corruption. Otherwise, consider pairs of
              different graph samples as negative pairs.
  
      :rtype: :class:`Tensor`
      """
      if zs_n is not None:
          assert len(zs_n) == len(zs)
          assert batch is not None
  
          jse = (
              JSE_local_global_negative_paired
              if neg_by_crpt
              else JSE_local_global
          )
  
          if len(zs) == 1:
              return jse(zs[0], zs_n[0], batch)
          elif len(zs) == 2:
              return jse(zs[0], zs_n[1], batch) + jse(
                  zs[1], zs_n[0], batch
              )
          else:
              assert len(zs) == len(sigma)
              loss = 0
              for i, j in itertools.combinations(range(len(zs)), 2):
                  if sigma[i][j]:
                      loss += jse(zs[i], zs_n[j], batch) + jse(
                          zs[j], zs_n[i], batch
                      )
              return loss
  
      else:
          jse = JSE_global_global
          if len(zs) == 2:
              return jse(zs[0], zs[1])
          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 += jse(zs[i], zs[j])
              return loss
  
file_path: DIG/dig/sslgraph/method/contrastive/objectives/jse.py
imports: import itertools
  import numpy
  import torch
  import torch.nn.functional
  import torch.nn
label_desc: Use image node_id 1 for calling a global function with example usage: JSE_loss(zs, zs_n, batch, sigma, neg_by_crpt) and returns: jse, jse, unknown, loss, loss