code: def JSE_local_global(z_g, z_n, batch):
      """
      Args:
          z_g: Tensor of shape [n_graphs, z_dim].
          z_n: Tensor of shape [n_nodes, z_dim].
          batch: Tensor of shape [n_graphs].
      """
      device = z_g.device
      num_graphs = z_g.shape[0]
      num_nodes = z_n.shape[0]
  
      pos_mask = torch.zeros((num_nodes, num_graphs)).to(device)
      neg_mask = torch.ones((num_nodes, num_graphs)).to(device)
      for nodeidx, graphidx in enumerate(batch):
          pos_mask[nodeidx][graphidx] = 1.0
          neg_mask[nodeidx][graphidx] = 0.0
  
      d_prime = torch.matmul(z_n, z_g.t())
  
      E_pos = get_expectation(d_prime * pos_mask, positive=True).sum()
      E_pos = E_pos / num_nodes
      E_neg = get_expectation(d_prime * neg_mask, positive=False).sum()
      E_neg = E_neg / (num_nodes * (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 3 for calling a global function with example usage: JSE_local_global(z_g, z_n, batch) and returns: unknown