code: def JSE_global_global(z1, z2):
      """
      Args:
          z1, z2: Tensor of shape [batch_size, z_dim].
      """
      device = z1.device
      num_graphs = z1.shape[0]
  
      pos_mask = torch.zeros((num_graphs, num_graphs)).to(device)
      neg_mask = torch.ones((num_graphs, num_graphs)).to(device)
      for graphidx in range(num_graphs):
          pos_mask[graphidx][graphidx] = 1.0
          neg_mask[graphidx][graphidx] = 0.0
  
      d_prime = torch.matmul(z1, z2.t())
  
      E_pos = get_expectation(d_prime * pos_mask, positive=True).sum()
      E_pos = E_pos / num_graphs
      E_neg = get_expectation(d_prime * neg_mask, positive=False).sum()
      E_neg = E_neg / (num_graphs * (num_graphs - 1))
      return E_neg - E_pos
  
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 4 for calling a global function with example usage: JSE_global_global(z1, z2) and returns: unknown